more type hints

glehmann · glehmann · commit 9fbbc6ed71ec · 2025-03-26T11:34:33.000+01:00
Signed-off-by: Gaëtan Lehmann &lt;gaetan.lehmann@vates.tech&gt;
diff --git a/lib/commands.py b/lib/commands.py
@@ -2,9 +2,8 @@
 import logging
 import shlex
 import subprocess
-from typing import Union
+from typing import List, Literal, overload, Union
 
-from _pytest.fixtures import _teardown_yield_fixture
 
 import lib.config as config
 from lib.netutil import wrap_ip
@@ -145,7 +144,34 @@ def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warning
 # The actual code is in _ssh().
 # This function is kept short for shorter pytest traces upon SSH failures, which are common,
 # as pytest prints the whole function definition that raised the SSHCommandFailed exception
-def ssh(hostname_or_ip, cmd, *, check=True, simple_output=True, suppress_fingerprint_warnings=True,
+@overload
+def ssh(hostname_or_ip: str, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[True] = True,
+        suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+        target_os: str = 'linux', decode: Literal[True] = True, options: List[str] = []) -> str:
+    ...
+@overload
+def ssh(hostname_or_ip: str, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[True] = True,
+        suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+        target_os: str = 'linux', decode: Literal[False], options: List[str] = []) -> bytes:
+    ...
+@overload
+def ssh(hostname_or_ip: str, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[False],
+        suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+        target_os: str = 'linux', decode: bool = True, options: List[str] = []) -> SSHResult:
+    ...
+@overload
+def ssh(hostname_or_ip: str, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[False],
+        suppress_fingerprint_warnings: bool = True, background: Literal[True],
+        target_os: str = 'linux', decode: bool = True, options: List[str] = []) -> subprocess.Popen:
+    ...
+@overload
+def ssh(hostname_or_ip: str, cmd: Union[str, List[str]], *, check=True, simple_output: bool = True,
+        suppress_fingerprint_warnings=True, background: bool = False,
+        target_os='linux', decode: bool = True, options: List[str] = []) \
+        -> Union[str, bytes, SSHResult, subprocess.Popen]:
+    ...
+def ssh(hostname_or_ip, cmd, *, check=True, simple_output=True,
+        suppress_fingerprint_warnings=True,
         background=False, target_os='linux', decode=True, options=[]):
     result_or_exc = _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
                          background, target_os, decode, options)
diff --git a/lib/common.py b/lib/common.py
@@ -5,9 +5,12 @@
 import time
 import traceback
 from enum import Enum
+from typing import TYPE_CHECKING, Dict, Optional, Union
 from uuid import UUID
 
 import lib.commands as commands
+if TYPE_CHECKING:
+    import lib.host
 
 class PackageManagerEnum(Enum):
     UNKNOWN = 1
@@ -154,9 +157,10 @@ def strtobool(str):
         return False
     raise ValueError("invalid truth value '{}'".format(str))
 
-def _param_get(host, xe_prefix, uuid, param_name, key=None, accept_unknown_key=False):
+def _param_get(host: 'lib.host.Host', xe_prefix: str, uuid: str, param_name: str, key: Optional[str] = None,
+               accept_unknown_key=False) -> Optional[str]:
     """ Common implementation for param_get. """
-    args = {'uuid': uuid, 'param-name': param_name}
+    args: Dict[str, Union[str, bool]] = {'uuid': uuid, 'param-name': param_name}
     if key is not None:
         args['param-key'] = key
     try:
diff --git a/lib/host.py b/lib/host.py
@@ -1,12 +1,12 @@
-import json
 import logging
 import os
 import shlex
+import subprocess
 import tempfile
 import uuid
 
 from packaging import version
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union, overload
 
 import lib.commands as commands
 import lib.pif as pif
@@ -40,9 +40,7 @@ class Host:
     def __init__(self, pool: 'lib.pool.Pool', hostname_or_ip):
         self.pool = pool
         self.hostname_or_ip = hostname_or_ip
-        self.inventory = None
-        self.uuid = None
-        self.xo_srv_id = None
+        self.xo_srv_id: Optional[str] = None
 
         h_data = host_data(self.hostname_or_ip)
         self.user = h_data['user']
@@ -59,11 +57,41 @@ def __init__(self, pool: 'lib.pool.Pool', hostname_or_ip):
     def __str__(self):
         return self.hostname_or_ip
 
-    def ssh(self, cmd, check=True, simple_output=True, suppress_fingerprint_warnings=True,
-            background=False, decode=True):
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[True] = True,
+            suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+            decode: Literal[True] = True) -> str:
+        ...
+
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[True] = True,
+            suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+            decode: Literal[False]) -> bytes:
+        ...
+
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[False],
+            suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+            decode: bool) -> commands.SSHResult:
+        ...
+
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: bool = True,
+            suppress_fingerprint_warnings: bool = True, background: Literal[True],
+            decode: bool = True) -> subprocess.Popen:
+        ...
+
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: bool = True,
+            suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True) \
+            -> Union[str, bytes, commands.SSHResult, subprocess.Popen]:
+        ...
+
+    def ssh(self, cmd, *, check=True, simple_output=True, suppress_fingerprint_warnings=True,
+            background=False, decode=True) -> Union[str, bytes, commands.SSHResult, subprocess.Popen]:
         return commands.ssh(self.hostname_or_ip, cmd, check=check, simple_output=simple_output,
-                            suppress_fingerprint_warnings=suppress_fingerprint_warnings, background=background,
-                            decode=decode)
+                            suppress_fingerprint_warnings=suppress_fingerprint_warnings,
+                            background=background, decode=decode)
 
     def ssh_with_result(self, cmd) -> commands.SSHResult:
         # doesn't raise if the command's return is nonzero, unless there's a SSH error
@@ -75,7 +103,18 @@ def scp(self, src, dest, check=True, suppress_fingerprint_warnings=True, local_d
             suppress_fingerprint_warnings=suppress_fingerprint_warnings, local_dest=local_dest
         )
 
-    def xe(self, action, args={}, check=True, simple_output=True, minimal=False, force=False):
+    @overload
+    def xe(self, action: str, args: Dict[str, Union[str, bool]] = {}, *, check: bool = True,
+           simple_output: Literal[True] = True, minimal: bool = False, force: bool = False) -> Union[bool, str]:
+        ...
+
+    @overload
+    def xe(self, action: str, args: Dict[str, Union[str, bool]] = {}, *, check: bool = True,
+           simple_output: Literal[False], minimal: bool = False, force: bool = False) -> commands.SSHResult:
+        ...
+
+    def xe(self, action, args={}, *, check=True, simple_output=True, minimal=False, force=False) \
+            -> Union[bool, str, commands.SSHResult]:
         maybe_param_minimal = ['--minimal'] if minimal else []
         maybe_param_force = ['--force'] if force else []
 
@@ -89,13 +128,14 @@ def stringify(key, value):
                 return ret.rstrip()
             return "{}={}".format(key, shlex.quote(value))
 
-        command = ['xe', action] + maybe_param_minimal + maybe_param_force + \
-                  [stringify(key, value) for key, value in args.items()]
+        command: List[str] = ['xe', action] + maybe_param_minimal + maybe_param_force + \
+                             [stringify(key, value) for key, value in args.items()]
         result = self.ssh(
             command,
             check=check,
             simple_output=simple_output
         )
+        assert isinstance(result, (str, commands.SSHResult))
 
         if result == 'true':
             return True
@@ -163,7 +203,7 @@ def execute_script(self, script_contents, shebang='sh', simple_output=True):
 
     def _get_xensource_inventory(self):
         output = self.ssh(['cat', '/etc/xensource-inventory'])
-        inventory = {}
+        inventory: dict[str, str] = {}
         for line in output.splitlines():
             key, raw_value = line.split('=')
             inventory[key] = raw_value.strip('\'')
@@ -202,7 +242,7 @@ def xo_server_add(self, username, password, label=None, unregister_first=True):
                 'allowUnauthorized': 'true',
                 'label': label
             }
-        )
+        ).decode()
         self.xo_srv_id = xo_srv_id
 
     def xo_server_status(self):
@@ -216,6 +256,7 @@ def xo_server_connected(self):
         return self.xo_server_status() == "connected"
 
     def xo_server_reconnect(self):
+        assert self.xo_srv_id is not None
         logging.info("Reconnect XO to host %s" % self)
         xo_cli('server.disable', {'id': self.xo_srv_id})
         xo_cli('server.enable', {'id': self.xo_srv_id})
@@ -285,6 +326,7 @@ def import_iso(self, uri, sr: SR):
             },
         )
 
+        download_path = None
         try:
             params = {'uuid': vdi_uuid}
             if '://' in uri:
@@ -293,7 +335,6 @@ def import_iso(self, uri, sr: SR):
                 self.ssh(f"curl -o '{download_path}' '{uri}'")
                 params['filename'] = download_path
             else:
-                download_path = None
                 params['filename'] = uri
             logging.info(f"Import ISO {uri}: name {random_name}, uuid {vdi_uuid}")
 
@@ -355,7 +396,7 @@ def get_last_yum_history_tid(self):
         """
         try:
             history_str = self.ssh(['yum', 'history', 'list', '--noplugins'])
-        except commands.SSHCommandFailed as e:
+        except commands.SSHCommandFailed:
             # yum history list fails if the list is empty, and it's also not possible to rollback
             # to before the first transaction, so "0" would not be appropriate as last transaction.
             # To workaround this, create transactions: install and remove a small package.
diff --git a/lib/pool.py b/lib/pool.py
@@ -1,6 +1,6 @@
 import logging
 import traceback
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, cast
 
 from packaging import version
 
@@ -30,7 +30,7 @@ def __init__(self, master_hostname_or_ip: str) -> None:
             if host_uuid != self.hosts[0].uuid:
                 host = Host(self, self.host_ip(host_uuid))
                 self.hosts.append(host)
-        self.uuid = self.master.xe('pool-list', minimal=True)
+        self.uuid = cast(str, self.master.xe('pool-list', minimal=True))
         self.saved_uefi_certs: Optional[Dict[str, Any]] = None
 
     def param_get(self, param_name, key=None, accept_unknown_key=False):
diff --git a/lib/pxe.py b/lib/pxe.py
@@ -1,4 +1,4 @@
-from lib.commands import ssh, scp, SSHCommandFailed
+from lib.commands import ssh, scp
 
 PXE_CONFIG_DIR = "/pxe/configs/custom"
 
diff --git a/lib/vm.py b/lib/vm.py
@@ -1,7 +1,8 @@
-import json
 import logging
 import os
+import subprocess
 import tempfile
+from typing import List, Literal, Optional, Union, overload
 
 import lib.commands as commands
 import lib.efi as efi
@@ -14,7 +15,7 @@
 class VM(BaseVM):
     def __init__(self, uuid, host):
         super().__init__(uuid, host)
-        self.ip = None
+        self.ip: Optional[str] = None
         self.previous_host = None # previous host when migrated or being migrated
         self.is_windows = self.param_get('platform', 'device_id', accept_unknown_key=True) == '0002'
         self.is_uefi = self.param_get('HVM-boot-params', 'firmware', accept_unknown_key=True) == 'uefi'
@@ -82,9 +83,31 @@ def try_get_and_store_ip(self):
             self.ip = ip
             return True
 
-    def ssh(self, cmd, check=True, simple_output=True, background=False, decode=True):
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[True] = True,
+            background: Literal[False] = False, decode: Literal[True] = True) -> str:
+        ...
+
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[True] = True,
+            background: Literal[False] = False, decode: Literal[False]) -> bytes:
+        ...
+
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: Literal[False],
+            background: Literal[False] = False, decode: bool = True) -> commands.SSHResult:
+        ...
+
+    @overload
+    def ssh(self, cmd: Union[str, List[str]], *, check: bool = True, simple_output: bool = True,
+            background: Literal[True], decode: bool = True) -> subprocess.Popen:
+        ...
+
+    def ssh(self, cmd: Union[str, List[str]], *, check=True, simple_output=True, background=False, decode=True) \
+            -> Union[str, bytes, commands.SSHResult, subprocess.Popen]:
         # raises by default for any nonzero return code
-        target_os = "windows" if self.is_windows else "linux"
+        assert self.ip is not None
+        target_os: str = "windows" if self.is_windows else "linux"
         return commands.ssh(self.ip, cmd, check=check, simple_output=simple_output, background=background,
                             target_os=target_os, decode=decode)
 
@@ -178,6 +201,7 @@ def exists(self):
         return self.host.pool_has_vm(self.uuid)
 
     def exists_on_previous_pool(self):
+        assert self.previous_host is not None
         return self.previous_host.pool_has_vm(self.uuid)
 
     def migrate(self, target_host, sr=None, network=None):
@@ -256,7 +280,7 @@ def get_residence_host(self):
         host_uuid = self.param_get('resident-on')
         return self.host.pool.get_host_by_uuid(host_uuid)
 
-    def start_background_process(self, cmd) -> str:
+    def start_background_process(self, cmd: str) -> str:
         script = "/tmp/bg_process.sh"
         pidfile = "/tmp/bg_process.pid"
         with tempfile.NamedTemporaryFile('w') as f:
@@ -420,7 +444,7 @@ def get_efi_var(self, var, guid):
         if not self.file_exists(efivarfs):
             return b''
 
-        data = self.ssh(['cat', efivarfs], simple_output=False, decode=False).stdout
+        data = self.ssh(['cat', efivarfs], decode=False)
 
         # The efivarfs file starts with the attributes, which are 4 bytes long
         return data[4:]
@@ -448,11 +472,10 @@ def get_all_efi_bins(self):
                 'do', 'echo', '$file', '$(head', '-c', magicsz, '$file);',
                 'done'
             ],
-            simple_output=False,
-            decode=False).stdout.split(b'\n')
+            decode=False).split(b'\n')
 
         magic = efi.EFI_HEADER_MAGIC.encode('ascii')
-        binaries = []
+        binaries: list[str] = []
         for f in files:
             if magic in f:
                 # Avoid decoding an unsplit f, as some headers are not utf8
@@ -590,6 +613,18 @@ def is_cert_present(self, key):
                             check=False, simple_output=False, decode=False)
         return res.returncode == 0
 
+    @overload
+    def execute_powershell_script(self, script_contents: str,
+                                  simple_output: Literal[True] = True,
+                                  prepend: str = "$ProgressPreference = 'SilentlyContinue';") -> str:
+        ...
+
+    @overload
+    def execute_powershell_script(self, script_contents: str,
+                                  simple_output: Literal[False],
+                                  prepend: str = "$ProgressPreference = 'SilentlyContinue';") -> commands.SSHResult:
+        ...
+
     def execute_powershell_script(
             self,
             script_contents: str,
diff --git a/lib/xo.py b/lib/xo.py
@@ -1,6 +1,24 @@
 import json
 import subprocess
+from typing import Any, Dict, Literal, Union, overload
 
+
+@overload
+def xo_cli(action: str, args: Dict[str, str] = {}, *, check: bool = True, simple_output: Literal[True] = True,
+           use_json: Literal[False] = False) -> bytes:
+    ...
+@overload
+def xo_cli(action: str, args: Dict[str, str] = {}, *, check: bool = True, simple_output: Literal[True] = True,
+           use_json: Literal[True]) -> Any:
+    ...
+@overload
+def xo_cli(action: str, args: Dict[str, str] = {}, *, check: bool = True, simple_output: Literal[False],
+           use_json: bool = False) -> subprocess.CompletedProcess:
+    ...
+@overload
+def xo_cli(action: str, args: Dict[str, str] = {}, *, check: bool = True, simple_output: bool = True,
+           use_json: bool = False) -> Union[subprocess.CompletedProcess, Any, bytes]:
+    ...
 def xo_cli(action, args={}, check=True, simple_output=True, use_json=False):
     run_array = ['xo-cli', action]
     if use_json:

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-from lib.commands import ssh, scp, SSHCommandFailed`
	`1`	`+from lib.commands import ssh, scp`
`2`	`2`
`3`	`3`	`PXE_CONFIG_DIR = "/pxe/configs/custom"`
`4`	`4`