Type testinfra;utils.ansible_runner

Author: philpep

testinfra/host.py

@@ -12,13 +12,15 @@
 
 import functools
 import os
+from typing import Any
 
 import testinfra.backend
+import testinfra.backend.base
 import testinfra.modules
 
 
 class Host:
-    _host_cache = {}  # type: ignore[var-annotated]
+    _host_cache: dict[tuple[str, frozenset[tuple[str, Any]]], "Host"] = {}
     _hosts_cache = {}  # type: ignore[var-annotated]
 
     def __init__(self, backend):
@@ -58,7 +60,9 @@ def find_command(self, command, extrapaths=("/sbin", "/usr/sbin")):
                 return path
         raise ValueError('cannot find "{}" command'.format(command))
 
-    def run(self, command, *args, **kwargs):
+    def run(
+        self, command: str, *args: str, **kwargs: Any
+    ) -> testinfra.backend.base.CommandResult:
         """Run given command and return rc (exit status), stdout and stderr
 
         >>> cmd = host.run("ls -l /etc/passwd")
@@ -85,9 +89,11 @@ def run(self, command, *args, **kwargs):
               'ls: cannot access /;echo inject: No such file or directory\\n'),
             command="ls -l '/;echo inject'")
         """
-        return self.backend.run(command, *args, **kwargs)
+        return self.backend.run(command, *args, **kwargs)  # type: ignore[no-any-return]
 
-    def run_expect(self, expected, command, *args, **kwargs):
+    def run_expect(
+        self, expected: list[int], command: str, *args: str, **kwargs: Any
+    ) -> testinfra.backend.base.CommandResult:
         """Run command and check it return an expected exit status
 
         :param expected: A list of expected exit status
@@ -105,7 +111,7 @@ def run_test(self, command, *args, **kwargs):
         """
         return self.run_expect([0, 1], command, *args, **kwargs)
 
-    def check_output(self, command, *args, **kwargs):
+    def check_output(self, command: str, *args: str, **kwargs: Any) -> str:
         """Get stdout of a command which has run successfully
 
         :returns: stdout without trailing newline
@@ -127,7 +133,7 @@ def __getattr__(self, name):
         )
 
     @classmethod
-    def get_host(cls, hostspec, **kwargs):
+    def get_host(cls, hostspec: str, **kwargs: Any) -> "Host":
         """Return a Host instance from `hostspec`
 
         `hostspec` should be like

testinfra/main.py

@@ -15,6 +15,6 @@
 import pytest
 
 
-def main():
+def main() -> int:
     warnings.warn("calling testinfra is deprecated, call py.test instead", stacklevel=1)
     return pytest.main()

testinfra/utils/ansible_runner.py

@@ -17,16 +17,17 @@
 import json
 import os
 import tempfile
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Iterator, Optional, Union
 
 import testinfra
+import testinfra.host
 
 __all__ = ["AnsibleRunner"]
 
 local = testinfra.get_host("local://")
 
 
-def get_ansible_config():
+def get_ansible_config() -> configparser.ConfigParser:
     fname = os.environ.get("ANSIBLE_CONFIG")
     if not fname:
         for possible in (
@@ -44,18 +45,29 @@ def get_ansible_config():
     return config
 
 
-def get_ansible_inventory(config, inventory_file):
+Inventory = dict[str, Any]
+
+
+def get_ansible_inventory(
+    config: configparser.ConfigParser, inventory_file: Optional[str]
+) -> Inventory:
     # Disable ansible verbosity to avoid
     # https://github.com/ansible/ansible/issues/59973
     cmd = "ANSIBLE_VERBOSITY=0 ansible-inventory --list"
     args = []
     if inventory_file:
         cmd += " -i %s"
         args += [inventory_file]
-    return json.loads(local.check_output(cmd, *args))
+    return json.loads(local.check_output(cmd, *args))  # type: ignore[no-any-return]
 
 
-def get_ansible_host(config, inventory, host, ssh_config=None, ssh_identity_file=None):
+def get_ansible_host(
+    config: configparser.ConfigParser,
+    inventory: Inventory,
+    host: str,
+    ssh_config: Optional[str] = None,
+    ssh_identity_file: Optional[str] = None,
+) -> Optional[testinfra.host.Host]:
     if is_empty_inventory(inventory):
         if host == "localhost":
             return testinfra.get_host("local://")
@@ -81,7 +93,7 @@ def get_ansible_host(config, inventory, host, ssh_config=None, ssh_identity_file
         "smart": "ssh",
     }.get(connection, connection)
 
-    options: Dict[str, Any] = {
+    options: dict[str, Any] = {
         "ansible_become": {
             "ini": {
                 "section": "privilege_escalation",
@@ -126,7 +138,9 @@ def get_ansible_host(config, inventory, host, ssh_config=None, ssh_identity_file
         },
     }
 
-    def get_config(name, default=None):
+    def get_config(
+        name: str, default: Union[None, bool, str] = None
+    ) -> Union[None, bool, str]:
         value = default
         option = options.get(name, {})
 
@@ -144,11 +158,12 @@ def get_config(name, default=None):
         return value
 
     testinfra_host = get_config("ansible_host", host)
+    assert isinstance(testinfra_host, str), testinfra_host
     user = get_config("ansible_user")
     password = get_config("ansible_ssh_pass")
     port = get_config("ansible_port")
 
-    kwargs: Dict[str, Union[str, bool]] = {}
+    kwargs: dict[str, Union[None, str, bool]] = {}
     if get_config("ansible_become", False):
         kwargs["sudo"] = True
     kwargs["sudo_user"] = get_config("ansible_become_user")
@@ -165,8 +180,8 @@ def get_config(name, default=None):
     kwargs["ssh_extra_args"] = " ".join(
         [
             config.get("ssh_connection", "ssh_args", fallback=""),
-            get_config("ansible_ssh_common_args", ""),
-            get_config("ansible_ssh_extra_args", ""),
+            get_config("ansible_ssh_common_args", ""),  # type: ignore[list-item]
+            get_config("ansible_ssh_extra_args", ""),  # type: ignore[list-item]
         ]
     ).strip()
 
@@ -191,20 +206,20 @@ def get_config(name, default=None):
     return testinfra.get_host(spec, **kwargs)
 
 
-def itergroup(inventory, group):
+def itergroup(inventory: Inventory, group: str) -> Iterator[str]:
     for host in inventory.get(group, {}).get("hosts", []):
         yield host
     for g in inventory.get(group, {}).get("children", []):
         for host in itergroup(inventory, g):
             yield host
 
 
-def is_empty_inventory(inventory):
+def is_empty_inventory(inventory: Inventory) -> bool:
     return not any(True for _ in itergroup(inventory, "all"))
 
 
 class AnsibleRunner:
-    _runners: Dict[Optional[str], "AnsibleRunner"] = {}
+    _runners: dict[Optional[str], "AnsibleRunner"] = {}
     _known_options = {
         # Boolean arguments.
         "become": {
@@ -243,12 +258,12 @@ class AnsibleRunner:
         },
     }
 
-    def __init__(self, inventory_file=None):
+    def __init__(self, inventory_file: Optional[str] = None):
         self.inventory_file = inventory_file
-        self._host_cache = {}
+        self._host_cache: dict[str, Optional[testinfra.host.Host]] = {}
         super().__init__()
 
-    def get_hosts(self, pattern="all"):
+    def get_hosts(self, pattern: str = "all") -> list[str]:
         inventory = self.inventory
         result = set()
         if is_empty_inventory(inventory):
@@ -271,18 +286,18 @@ def get_hosts(self, pattern="all"):
         return sorted(result)
 
     @functools.cached_property
-    def inventory(self):
+    def inventory(self) -> Inventory:
         return get_ansible_inventory(self.ansible_config, self.inventory_file)
 
     @functools.cached_property
-    def ansible_config(self):
+    def ansible_config(self) -> configparser.ConfigParser:
         return get_ansible_config()
 
-    def get_variables(self, host):
+    def get_variables(self, host: str) -> dict[str, Any]:
         inventory = self.inventory
         # inventory_hostname, group_names and groups are for backward
         # compatibility with testinfra 2.X
-        hostvars = inventory["_meta"].get("hostvars", {}).get(host, {})
+        hostvars: dict[str, Any] = inventory["_meta"].get("hostvars", {}).get(host, {})
         hostvars.setdefault("inventory_hostname", host)
         group_names = []
         groups = {}
@@ -296,7 +311,7 @@ def get_variables(self, host):
         hostvars.setdefault("groups", groups)
         return hostvars
 
-    def get_host(self, host, **kwargs):
+    def get_host(self, host: str, **kwargs: Any) -> Optional[testinfra.host.Host]:
         try:
             return self._host_cache[host]
         except KeyError:
@@ -305,14 +320,14 @@ def get_host(self, host, **kwargs):
             )
             return self._host_cache[host]
 
-    def options_to_cli(self, options):
+    def options_to_cli(self, options: dict[str, Any]) -> tuple[str, list[str]]:
         verbose = options.pop("verbose", 0)
 
         args = {"become": False, "check": True}
         args.update(options)
 
-        cli: List[str] = []
-        cli_args: List[str] = []
+        cli: list[str] = []
+        cli_args: list[str] = []
         if verbose:
             cli.append("-" + "v" * verbose)
         for arg_name, value in args.items():
@@ -334,7 +349,14 @@ def options_to_cli(self, options):
                 raise TypeError("Unsupported argument type '{}'.".format(opt_type))
         return " ".join(cli), cli_args
 
-    def run_module(self, host, module_name, module_args, get_encoding=None, **options):
+    def run_module(
+        self,
+        host: str,
+        module_name: str,
+        module_args: str,
+        get_encoding: Optional[Callable[[], str]] = None,
+        **options: Any,
+    ) -> Any:
         cmd, args = "ansible --tree %s", []
         if self.inventory_file:
             cmd += " -i %s"
@@ -375,7 +397,7 @@ def run_module(self, host, module_name, module_args, get_encoding=None, **option
                     return json.load(f)
 
     @classmethod
-    def get_runner(cls, inventory):
+    def get_runner(cls, inventory: Optional[str]) -> "AnsibleRunner":
         try:
             return cls._runners[inventory]
         except KeyError: