feat: add get_hosts_by_ansible function

Author: amarao

doc/source/backends.rst

@@ -194,6 +194,51 @@ https://docs.ansible.com/ansible/latest/reference_appendices/config.html
 * ``ANSIBLE_BECOME_USER``
 * ``ANSIBLE_BECOME``
 
+Advanced hosts expressions for ansible
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+It is possible to use most of the
+`Ansible host expressions <https://docs.ansible.com/ansible/latest/inventory_guide/intro_patterns.html>`
+in Testinfra.
+
+Supported:
+
+* ``&``, ``!``, ``,``, ``:``
+* glob expressions (``server*`` with match server1, server2, etc)
+* ranges (``[x:y]``) are supported with replacement with round brackets.
+  ``mygroup[1:2]`` should be written as ``mygroup(1:2)``, this is due to limitation
+  of what is allowed in the host part of the URL.
+  When host expression is passed to ansible for parsing, ``()`` are replaced with
+  ``[]``.
+
+Regular expressions (starting with '~') are not supported due to limitations
+for allowed characters in the host part of the URL.
+
+When testinfra parses host expressions, it choose:
+
+* A simple resolver, if there is no host expression (e.g. a single group,
+  hostname, or glob pattern)
+* Ansible resolver, which covers most cases. It requires to have ansible-core
+  been present on the controller (host, where pytest is running). It imports
+  part of ansible to do expression evaluation, os it's slower.
+
+Examples of the simple host expression (Ansible is not used for parsing):
+
+* ``ansible://debian_bookworm``
+* ``ansible://user@debian_bookworm?force_ansible=True&sudo=True``
+* ``ansible://host*``
+
+Examples of the Ansible-parsed host expressions:
+
+* ``ansible://group1,!group3`` (hosts in group1 but not in group3)
+* ``ansible://group1(0)`` (the first host in the group). This can be used as a substitute
+  for run_once.
+* ``ansible://group1,&group3`` (hosts in both group1 and group2)
+* ``ansible://group1,group2,!group3,example*`` (hosts in group1 or group2 but not
+  in group3, and hosts matching regular expression ``(example1.*)``)
+* ``ansible://group1,group2,!group3,example*?force_ansible=True&sudo=True``
+  (the same, but forcing Ansible backend and adds sudo)
+
 kubectl
 ~~~~~~~
 

mypy.ini

@@ -19,3 +19,6 @@ ignore_missing_imports = True
 
 [mypy-setuptools_scm.*]
 ignore_missing_imports = True
+
+[mypy-ansible.*]
+ignore_missing_imports = True

test/test_ansible_host_expressions.py

@@ -0,0 +1,99 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tempfile
+
+import pytest
+
+from testinfra.utils.ansible_runner import AnsibleRunner
+
+INVENTORY = b"""
+---
+group1:
+  hosts:
+    host1:
+    host2:
+group2:
+  hosts:
+    host3:
+    host4:
+group3:
+  hosts:
+    host1:
+    host4:
+group4:
+   hosts:
+    example1:
+    example11:
+    example2:
+"""
+
+
+@pytest.fixture(scope="module")
+def get_hosts():
+    with tempfile.NamedTemporaryFile() as f:
+        f.write(INVENTORY)
+        f.flush()
+
+        def _get_hosts(spec):
+            return AnsibleRunner(f.name).get_hosts(spec)
+
+        yield _get_hosts
+
+
+@pytest.fixture(scope="function")
+def get_env_var():
+    old_value = os.environ.get("ANSIBLE_INVENTORY")
+    with tempfile.NamedTemporaryFile() as f:
+        f.write(INVENTORY)
+        f.flush()
+        os.environ["ANSIBLE_INVENTORY"] = f.name
+
+        def _get_env_hosts(spec):
+            return AnsibleRunner(None).get_hosts(spec)
+
+        yield _get_env_hosts
+    if old_value:
+        os.environ["ANSIBLE_INVENTORY"] = old_value
+    else:
+        del os.environ["ANSIBLE_INVENTORY"]
+
+
+def test_ansible_host_expressions_index(get_hosts):
+    assert get_hosts("group1(0)") == ["host1"]
+
+
+def test_ansible_host_expressions_negative_index(get_hosts):
+    assert get_hosts("group1(-1)") == ["host2"]
+
+
+def test_ansible_host_expressions_not(get_hosts):
+    assert get_hosts("group1,!group3") == ["host2"]
+
+
+def test_ansible_host_expressions_and(get_hosts):
+    assert get_hosts("group1,&group3") == ["host1"]
+
+
+def test_ansible_host_complicated_expression(get_hosts):
+    expression = "group1,group2,!group3,example1*"
+    assert set(get_hosts(expression)) == {"host2", "host3", "example1", "example11"}
+
+
+def test_ansible_host_regexp(get_hosts):
+    with pytest.raises(ValueError):
+        get_hosts("~example1*")
+
+
+def test_ansible_host_with_ansible_inventory_env_var(get_env_var):
+    assert set(get_env_var("host1,example1*")) == {"host1", "example1", "example11"}

testinfra/utils/ansible_runner.py

@@ -277,6 +277,42 @@ def __init__(self, inventory_file: Optional[str] = None):
         self._host_cache: dict[str, Optional[testinfra.host.Host]] = {}
         super().__init__()
 
+    def get_hosts_by_ansible(self, host_expression: str) -> list[str]:
+        """Evaluate ansible host expression to get host list.
+
+        Example of such expression:
+
+        'foo,&bar,!baz[2:],foobar[-3:-4],foofoo-*,~someth.+'
+
+        See https://docs.ansible.com/ansible/latest/inventory_guide/intro_patterns.html#common-patterns
+        """
+        from ansible.inventory.manager import InventoryManager
+        from ansible.parsing.dataloader import DataLoader
+
+        # We can't support 'group[1:2]', expressions 'as is',
+        # because urllib from python 3.13+ rejects hostnames with invalid
+        #  IPv6 addresses (in square brakets)
+        # E   ValueError: Invalid IPv6 URL
+        # We ask user to use round brakets in testinfra 'URL', and
+        # replace it back to ansible-compatible expression here.
+        host_expression = host_expression.replace("(", "[")
+        host_expression = host_expression.replace(")", "]")
+        if self.inventory_file:
+            sources = [self.inventory_file]
+        else:
+            # we search for other options only if inventory is not passed
+            # explicitely.
+            # Inside ansible, 'ANSIBLE_INVENTORY' is 'DEFAULT_HOST_LIST'
+            # We respect both ANSIBLE_INVENTORY env var and ansible.cfg
+            from ansible.config.manager import ConfigManager
+
+            sources = ConfigManager().get_config_value("DEFAULT_HOST_LIST")
+
+        loader = DataLoader()
+        inv = InventoryManager(loader=loader, sources=sources)
+        hosts = [h.name for h in inv.list_hosts(host_expression)]
+        return list(hosts)
+
     def get_hosts(self, pattern: str = "all") -> list[str]:
         inventory = self.inventory
         result = set()
@@ -290,13 +326,22 @@ def get_hosts(self, pattern: str = "all") -> list[str]:
                     "only implicit localhost is available"
                 )
         else:
-            for group in inventory:
-                groupmatch = fnmatch.fnmatch(group, pattern)
-                if groupmatch:
-                    result |= set(itergroup(inventory, group))
-                for host in inventory[group].get("hosts", []):
-                    if fnmatch.fnmatch(host, pattern):
-                        result.add(host)
+            if "~" in pattern:
+                raise ValueError(
+                    "Regular expressions are not supported in host expression. "
+                    "Found '~' in the host expression."
+                )
+            special_char_list = "!&,:()"  # signs of host expression
+            if any(ch in special_char_list for ch in pattern):
+                result = result.union(self.get_hosts_by_ansible(pattern))
+            else:
+                for group in inventory:
+                    groupmatch = fnmatch.fnmatch(group, pattern)
+                    if groupmatch:
+                        result |= set(itergroup(inventory, group))
+                    for host in inventory[group].get("hosts", []):
+                        if fnmatch.fnmatch(host, pattern):
+                            result.add(host)
         return sorted(result)
 
     @functools.cached_property

tox.ini

@@ -36,6 +36,7 @@ commands=
 description = Performs typing check
 extras =
   typing
+  ansible
 usedevelop=True
 commands=
   mypy