Type testinfra.host

Author: philpep

testinfra/backend/__init__.py

@@ -13,6 +13,10 @@
 import importlib
 import os
 import urllib.parse
+from typing import TYPE_CHECKING, Any, Iterable
+
+if TYPE_CHECKING:
+    import testinfra.backend.base
 
 BACKENDS = {
     "local": "testinfra.backend.local.LocalBackend",
@@ -31,13 +35,13 @@
 }
 
 
-def get_backend_class(connection):
+def get_backend_class(connection: str) -> type["testinfra.backend.base.BaseBackend"]:
     try:
         classpath = BACKENDS[connection]
     except KeyError:
         raise RuntimeError("Unknown connection type '{}'".format(connection))
     module, name = classpath.rsplit(".", 1)
-    return getattr(importlib.import_module(module), name)
+    return getattr(importlib.import_module(module), name)  # type: ignore[no-any-return]
 
 
 def parse_hostspec(hostspec):
@@ -75,7 +79,7 @@ def parse_hostspec(hostspec):
     return host, kw
 
 
-def get_backend(hostspec, **kwargs):
+def get_backend(hostspec: str, **kwargs: Any) -> "testinfra.backend.base.BaseBackend":
     host, kw = parse_hostspec(hostspec)
     for k, v in kwargs.items():
         kw.setdefault(k, v)
@@ -86,7 +90,9 @@ def get_backend(hostspec, **kwargs):
     return klass(host, **kw)
 
 
-def get_backends(hosts, **kwargs):
+def get_backends(
+    hosts: Iterable[str], **kwargs: Any
+) -> list["testinfra.backend.base.BaseBackend"]:
     backends = {}
     for hostspec in hosts:
         host, kw = parse_hostspec(hostspec)

testinfra/backend/base.py

@@ -17,7 +17,10 @@
 import shlex
 import subprocess
 import urllib.parse
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    import testinfra.host
 
 logger = logging.getLogger("testinfra")
 
@@ -117,29 +120,29 @@ def __init__(
         self,
         hostname: str,
         sudo: bool = False,
-        sudo_user: Optional[bool] = None,
+        sudo_user: Optional[str] = None,
         *args: Any,
         **kwargs: Any,
     ):
-        self._encoding = None
-        self._host = None
+        self._encoding: Optional[str] = None
+        self._host: Optional["testinfra.host.Host"] = None
         self.hostname = hostname
         self.sudo = sudo
         self.sudo_user = sudo_user
         super().__init__()
 
-    def set_host(self, host):
+    def set_host(self, host: "testinfra.host.Host") -> None:
         self._host = host
 
     @classmethod
-    def get_connection_type(cls):
+    def get_connection_type(cls) -> str:
         """Return the connection backend used as string.
 
         Can be local, paramiko, ssh, docker, salt or ansible
         """
         return cls.NAME
 
-    def get_hostname(self):
+    def get_hostname(self) -> str:
         """Return the hostname (for testinfra) of the remote or local system
 
 
@@ -166,11 +169,11 @@ def test(TestinfraBackend):
         """
         return self.hostname
 
-    def get_pytest_id(self):
+    def get_pytest_id(self) -> str:
         return self.get_connection_type() + "://" + self.get_hostname()
 
     @classmethod
-    def get_hosts(cls, host, **kwargs):
+    def get_hosts(cls, host: str, **kwargs: Any) -> list[str]:
         if host is None:
             raise RuntimeError(
                 "One or more hosts is required with the {} backend".format(
@@ -180,41 +183,41 @@ def get_hosts(cls, host, **kwargs):
         return [host]
 
     @staticmethod
-    def quote(command, *args):
+    def quote(command: str, *args: str) -> str:
         if args:
             return command % tuple(shlex.quote(a) for a in args)  # noqa: S001
         return command
 
-    def get_sudo_command(self, command, sudo_user):
+    def get_sudo_command(self, command: str, sudo_user: Optional[str]) -> str:
         if sudo_user is None:
             return self.quote("sudo /bin/sh -c %s", command)
         return self.quote("sudo -u %s /bin/sh -c %s", sudo_user, command)
 
-    def get_command(self, command, *args):
+    def get_command(self, command: str, *args: str) -> str:
         command = self.quote(command, *args)
         if self.sudo:
             command = self.get_sudo_command(command, self.sudo_user)
         return command
 
-    def run(self, command, *args, **kwargs):
+    def run(self, command: str, *args: str, **kwargs: Any) -> CommandResult:
         raise NotImplementedError
 
-    def run_local(self, command, *args):
+    def run_local(self, command: str, *args: str) -> CommandResult:
         command = self.quote(command, *args)
-        command = self.encode(command)
+        cmd = self.encode(command)
         p = subprocess.Popen(
-            command,
+            cmd,
             shell=True,
             stdin=subprocess.PIPE,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
         )
         stdout, stderr = p.communicate()
-        result = self.result(p.returncode, command, stdout, stderr)
+        result = self.result(p.returncode, cmd, stdout, stderr)
         return result
 
     @staticmethod
-    def parse_hostspec(hostspec):
+    def parse_hostspec(hostspec: str) -> HostSpec:
         name = hostspec
         port = None
         user = None
@@ -246,14 +249,14 @@ def parse_hostspec(hostspec):
         return HostSpec(name, port, user, password)
 
     @staticmethod
-    def parse_containerspec(containerspec):
+    def parse_containerspec(containerspec: str) -> tuple[str, Optional[str]]:
         name = containerspec
         user = None
         if "@" in name:
             user, name = name.split("@", 1)
         return name, user
 
-    def get_encoding(self):
+    def get_encoding(self) -> str:
         encoding = None
         for python in ("python3", "python"):
             cmd = self.run(
@@ -274,7 +277,7 @@ def get_encoding(self):
         return encoding
 
     @property
-    def encoding(self):
+    def encoding(self) -> str:
         if self._encoding is None:
             self._encoding = self.get_encoding()
         return self._encoding

testinfra/host.py

@@ -12,38 +12,44 @@
 
 import functools
 import os
+from collections.abc import Iterable
 from typing import Any
 
 import testinfra.backend
 import testinfra.backend.base
 import testinfra.modules
+import testinfra.modules.base
 
 
 class Host:
     _host_cache: dict[tuple[str, frozenset[tuple[str, Any]]], "Host"] = {}
-    _hosts_cache = {}  # type: ignore[var-annotated]
+    _hosts_cache: dict[
+        tuple[frozenset[str], frozenset[tuple[str, Any]]], list["Host"]
+    ] = {}
 
-    def __init__(self, backend):
+    def __init__(self, backend: testinfra.backend.base.BaseBackend):
         self.backend = backend
         super().__init__()
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "<testinfra.host.Host {}>".format(self.backend.get_pytest_id())
 
     @functools.cached_property
-    def has_command_v(self):
+    def has_command_v(self) -> bool:
         """Return True if `command -v` is available"""
         return self.run("command -v command").rc == 0
 
-    def exists(self, command):
+    def exists(self, command: str) -> bool:
         """Return True if given command exist in $PATH"""
         if self.has_command_v:
             out = self.run("command -v %s", command)
         else:
             out = self.run_expect([0, 1], "which %s", command)
         return out.rc == 0
 
-    def find_command(self, command, extrapaths=("/sbin", "/usr/sbin")):
+    def find_command(
+        self, command: str, extrapaths: Iterable[str] = ("/sbin", "/usr/sbin")
+    ) -> str:
         """Return path of given command
 
         raise ValueError if command cannot be found
@@ -89,7 +95,7 @@ def run(
               'ls: cannot access /;echo inject: No such file or directory\\n'),
             command="ls -l '/;echo inject'")
         """
-        return self.backend.run(command, *args, **kwargs)  # type: ignore[no-any-return]
+        return self.backend.run(command, *args, **kwargs)
 
     def run_expect(
         self, expected: list[int], command: str, *args: str, **kwargs: Any
@@ -104,7 +110,9 @@ def run_expect(
         assert out.rc in expected, "Unexpected exit code {} for {}".format(out.rc, out)
         return out
 
-    def run_test(self, command, *args, **kwargs):
+    def run_test(
+        self, command: str, *args: str, **kwargs: Any
+    ) -> testinfra.backend.base.CommandResult:
         """Run command and check it return an exit status of 0 or 1
 
         :raises: AssertionError
@@ -122,7 +130,7 @@ def check_output(self, command: str, *args: str, **kwargs: Any) -> str:
         assert out.rc == 0, "Unexpected exit code {} for {}".format(out.rc, out)
         return out.stdout.rstrip("\r\n")
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> type[testinfra.modules.base.Module]:
         if name in testinfra.modules.modules:
             module_class = testinfra.modules.get_module_class(name)
             obj = module_class.get_module(self)
@@ -157,7 +165,7 @@ def get_host(cls, hostspec: str, **kwargs: Any) -> "Host":
         return cache[key]
 
     @classmethod
-    def get_hosts(cls, hosts, **kwargs):
+    def get_hosts(cls, hosts: Iterable[str], **kwargs: Any) -> list["Host"]:
         key = (frozenset(hosts), frozenset(kwargs.items()))
         cache = cls._hosts_cache
         if key not in cache:

testinfra/modules/__init__.py

@@ -11,6 +11,10 @@
 # limitations under the License.
 
 import importlib
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import testinfra.modules.base
 
 modules = {
     "addr": "addr:Addr",
@@ -42,8 +46,8 @@
 }
 
 
-def get_module_class(name):
+def get_module_class(name: str) -> type["testinfra.modules.base.Module"]:
     modname, classname = modules[name].split(":")
     modname = ".".join([__name__, modname])
     module = importlib.import_module(modname)
-    return getattr(module, classname)
+    return getattr(module, classname)  # type: ignore[no-any-return]

testinfra/modules/base.py

@@ -9,17 +9,17 @@
 # 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 typing
+from typing import TYPE_CHECKING
 
 
 class Module:
-    if typing.TYPE_CHECKING:
+    if TYPE_CHECKING:
         import testinfra.host
 
         _host: testinfra.host.Host
 
     @classmethod
-    def get_module(cls, _host):
+    def get_module(cls, _host: "testinfra.host.Host") -> type["Module"]:
         klass = cls.get_module_class(_host)
         return type(
             klass.__name__,