Support for Python 3.11+ tomllib for inventory (#77435)

lib/ansible/cli/inventory.py

@@ -180,13 +180,13 @@ class InventoryCLI(CLI):
             from ansible.parsing.yaml.dumper import AnsibleDumper
             results = to_text(yaml.dump(stuff, Dumper=AnsibleDumper, default_flow_style=False, allow_unicode=True))
         elif context.CLIARGS['toml']:
-            from ansible.plugins.inventory.toml import toml_dumps, HAS_TOML
-            if not HAS_TOML:
-                raise AnsibleError(
-                    'The python "toml" library is required when using the TOML output format'
-                )
+            from ansible.plugins.inventory.toml import toml_dumps
             try:
                 results = toml_dumps(stuff)
+            except TypeError as e:
+                raise AnsibleError(
+                    'The source inventory contains a value that cannot be represented in TOML: %s' % e
+                )
             except KeyError as e:
                 raise AnsibleError(
                     'The source inventory contains a non-string key (%s) which cannot be represented in TOML. '

lib/ansible/plugins/inventory/toml.py

@@ -12,7 +12,8 @@ DOCUMENTATION = r'''
         - TOML based inventory format
         - File MUST have a valid '.toml' file extension
     notes:
-        - Requires the 'toml' python library
+        - >
+          Requires one of the following python libraries: 'toml', 'tomli', or 'tomllib'
 '''
 
 EXAMPLES = r'''# fmt: toml
@@ -92,7 +93,7 @@ import typing as t
 from collections.abc import MutableMapping, MutableSequence
 from functools import partial
 
-from ansible.errors import AnsibleFileNotFound, AnsibleParserError
+from ansible.errors import AnsibleFileNotFound, AnsibleParserError, AnsibleRuntimeError
 from ansible.module_utils._text import to_bytes, to_native, to_text
 from ansible.module_utils.six import string_types, text_type
 from ansible.parsing.yaml.objects import AnsibleSequence, AnsibleUnicode
@@ -100,16 +101,37 @@ from ansible.plugins.inventory import BaseFileInventoryPlugin
 from ansible.utils.display import Display
 from ansible.utils.unsafe_proxy import AnsibleUnsafeBytes, AnsibleUnsafeText
 
+HAS_TOML = False
 try:
     import toml
     HAS_TOML = True
 except ImportError:
-    HAS_TOML = False
+    pass
+
+HAS_TOMLIW = False
+try:
+    import tomli_w  # type: ignore[import]
+    HAS_TOMLIW = True
+except ImportError:
+    pass
+
+HAS_TOMLLIB = False
+try:
+    import tomllib  # type: ignore[import]
+    HAS_TOMLLIB = True
+except ImportError:
+    try:
+        import tomli as tomllib  # type: ignore[no-redef]
+        HAS_TOMLLIB = True
+    except ImportError:
+        pass
 
 display = Display()
 
 
+# dumps
 if HAS_TOML and hasattr(toml, 'TomlEncoder'):
+    # toml>=0.10.0
     class AnsibleTomlEncoder(toml.TomlEncoder):
         def __init__(self, *args, **kwargs):
             super(AnsibleTomlEncoder, self).__init__(*args, **kwargs)
@@ -122,20 +144,39 @@ if HAS_TOML and hasattr(toml, 'TomlEncoder'):
             })
     toml_dumps = partial(toml.dumps, encoder=AnsibleTomlEncoder())  # type: t.Callable[[t.Any], str]
 else:
+    # toml<0.10.0
+    # tomli-w
     def toml_dumps(data):  # type: (t.Any) -> str
+        if HAS_TOML:
             return toml.dumps(convert_yaml_objects_to_native(data))
+        elif HAS_TOMLIW:
+            return tomli_w.dumps(convert_yaml_objects_to_native(data))
+        raise AnsibleRuntimeError(
+            'The python "toml" or "tomli-w" library is required when using the TOML output format'
+        )
+
+# loads
+if HAS_TOML:
+    # prefer toml if installed, since it supports both encoding and decoding
+    toml_loads = toml.loads  # type: ignore[assignment]
+    TOMLDecodeError = toml.TomlDecodeError  # type: t.Any
+elif HAS_TOMLLIB:
+    toml_loads = tomllib.loads  # type: ignore[assignment]
+    TOMLDecodeError = tomllib.TOMLDecodeError  # type: t.Any  # type: ignore[no-redef]
 
 
 def convert_yaml_objects_to_native(obj):
-    """Older versions of the ``toml`` python library, don't have a pluggable
-    way to tell the encoder about custom types, so we need to ensure objects
-    that we pass are native types.
+    """Older versions of the ``toml`` python library, and tomllib, don't have
+    a pluggable way to tell the encoder about custom types, so we need to
+    ensure objects that we pass are native types.
 
-    Only used on ``toml<0.10.0`` where ``toml.TomlEncoder`` is missing.
+    Used with:
+      - ``toml<0.10.0`` where ``toml.TomlEncoder`` is missing
+      - ``tomli`` or ``tomllib``
 
     This function recurses an object and ensures we cast any of the types from
     ``ansible.parsing.yaml.objects`` into their native types, effectively cleansing
-    the data before we hand it over to ``toml``
+    the data before we hand it over to the toml library.
 
     This function doesn't directly check for the types from ``ansible.parsing.yaml.objects``
     but instead checks for the types those objects inherit from, to offer more flexibility.
@@ -207,8 +248,8 @@ class InventoryModule(BaseFileInventoryPlugin):
 
         try:
             (b_data, private) = self.loader._get_file_contents(file_name)
-            return toml.loads(to_text(b_data, errors='surrogate_or_strict'))
-        except toml.TomlDecodeError as e:
+            return toml_loads(to_text(b_data, errors='surrogate_or_strict'))
+        except TOMLDecodeError as e:
             raise AnsibleParserError(
                 'TOML file (%s) is invalid: %s' % (file_name, to_native(e)),
                 orig_exc=e
@@ -226,9 +267,11 @@ class InventoryModule(BaseFileInventoryPlugin):
 
     def parse(self, inventory, loader, path, cache=True):
         ''' parses the inventory file '''
-        if not HAS_TOML:
+        if not HAS_TOMLLIB and not HAS_TOML:
+            # tomllib works here too, but we don't call it out in the error,
+            # since you either have it or not as part of cpython stdlib >= 3.11
             raise AnsibleParserError(
-                'The TOML inventory plugin requires the python "toml" library'
+                'The TOML inventory plugin requires the python "toml", or "tomli" library'
             )
 
         super(InventoryModule, self).parse(inventory, loader, path)

test/integration/targets/ansible-inventory/files/valid_sample.toml

@@ -0,0 +1,2 @@
+[somegroup.hosts.something]
+foo = "bar"

test/integration/targets/ansible-inventory/runme.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+
+source virtualenv.sh
+export ANSIBLE_ROLES_PATH=../
+set -euvx
+
+ansible-playbook test.yml "$@"

test/integration/targets/ansible-inventory/tasks/main.yml

@@ -82,30 +82,6 @@
         - result is failed
         - '"ERROR! Could not match supplied host pattern, ignoring: invalid" in result.stderr'
 
-- name: Install toml package
-  pip:
-    name:
-      - toml
-    state: present
-
-- name: "test option: --toml with valid group name"
-  command: ansible-inventory --list --toml -i {{ role_path }}/files/valid_sample.yml
-  register: result
-
-- assert:
-    that:
-        - result is succeeded
-
-- name: "test option: --toml with invalid group name"
-  command: ansible-inventory --list --toml -i {{ role_path }}/files/invalid_sample.yml
-  ignore_errors: true
-  register: result
-
-- assert:
-    that:
-        - result is failed
-        - '"ERROR! The source inventory contains a non-string key" in result.stderr'
-
 - name: "test json output with unicode characters"
   command: ansible-inventory --list -i {{ role_path }}/files/unicode.yml
   register: result
@@ -154,28 +130,18 @@
         name: unicode_inventory.yaml
         state: absent
 
-- block:
-  - name: "test toml output with unicode characters"
-    command: ansible-inventory --list --toml -i {{ role_path }}/files/unicode.yml
-    register: result
-
-  - assert:
-      that:
-          - result is succeeded
-          - result.stdout is contains('příbor')
-
-  - block:
-    - name: "test toml output file with unicode characters"
-      command: ansible-inventory --list --toml --output unicode_inventory.toml -i {{ role_path }}/files/unicode.yml
-
-    - set_fact:
-        toml_inventory_file: "{{ lookup('file', 'unicode_inventory.toml') | string }}"
-
-    - assert:
-        that:
-            - toml_inventory_file is contains('příbor')
-    always:
-      - file:
-          name: unicode_inventory.toml
-          state: absent
-  when: ansible_python.version.major|int == 3
+- include_tasks: toml.yml
+  loop:
+    -
+      - toml<0.10.0
+    -
+      - toml
+    -
+      - tomli
+      - tomli-w
+    -
+      - tomllib
+      - tomli-w
+  loop_control:
+    loop_var: toml_package
+  when: toml_package is not contains 'tomllib' or (toml_package is contains 'tomllib' and ansible_facts.python.version_info >= [3, 11])

test/integration/targets/ansible-inventory/tasks/toml.yml

@@ -0,0 +1,66 @@
+- name: Ensure no toml packages are installed
+  pip:
+    name:
+      - tomli
+      - tomli-w
+      - toml
+    state: absent
+
+- name: Install toml package
+  pip:
+    name: '{{ toml_package|difference(["tomllib"]) }}'
+    state: present
+
+- name: test toml parsing
+  command: ansible-inventory --list --toml -i {{ role_path }}/files/valid_sample.toml
+  register: toml_in
+
+- assert:
+    that:
+      - >
+        'foo = "bar"' in toml_in.stdout
+
+- name: "test option: --toml with valid group name"
+  command: ansible-inventory --list --toml -i {{ role_path }}/files/valid_sample.yml
+  register: result
+
+- assert:
+    that:
+        - result is succeeded
+
+- name: "test option: --toml with invalid group name"
+  command: ansible-inventory --list --toml -i {{ role_path }}/files/invalid_sample.yml
+  ignore_errors: true
+  register: result
+
+- assert:
+    that:
+        - result is failed
+        - >
+          "ERROR! The source inventory contains" in result.stderr
+
+- block:
+  - name: "test toml output with unicode characters"
+    command: ansible-inventory --list --toml -i {{ role_path }}/files/unicode.yml
+    register: result
+
+  - assert:
+      that:
+          - result is succeeded
+          - result.stdout is contains('příbor')
+
+  - block:
+    - name: "test toml output file with unicode characters"
+      command: ansible-inventory --list --toml --output unicode_inventory.toml -i {{ role_path }}/files/unicode.yml
+
+    - set_fact:
+        toml_inventory_file: "{{ lookup('file', 'unicode_inventory.toml') | string }}"
+
+    - assert:
+        that:
+            - toml_inventory_file is contains('příbor')
+    always:
+      - file:
+          name: unicode_inventory.toml
+          state: absent
+  when: ansible_python.version.major|int == 3

test/integration/targets/ansible-inventory/test.yml

@@ -0,0 +1,3 @@
+- hosts: localhost
+  roles:
+    - ansible-inventory