refactor: make BaseCmdResult and subclasses generic to eliminate runtime type narrowing

lib/commands.py

@@ -9,7 +9,7 @@
 import lib.config as config
 from lib.netutil import wrap_ip
 
-from typing import TYPE_CHECKING, List, Literal, overload
+from typing import TYPE_CHECKING, Generic, List, Literal, TypeVar, overload
 
 if TYPE_CHECKING:
     from lib.common import HostAddress
@@ -39,19 +39,21 @@ def __init__(self, returncode: int, stdout: str, cmd: str | list[str]):
             f'Local command ({cmd}) failed with return code {returncode}{msg_end}'
         )
 
-class BaseCmdResult:
+ResultOutputT = TypeVar('ResultOutputT', str, bytes)
+
+class BaseCmdResult(Generic[ResultOutputT]):
     __slots__ = 'returncode', 'stdout'
 
-    def __init__(self, returncode: int, stdout: str | bytes):
+    def __init__(self, returncode: int, stdout: ResultOutputT):
         self.returncode = returncode
-        self.stdout = stdout
+        self.stdout: ResultOutputT = stdout
 
-class SSHResult(BaseCmdResult):
-    def __init__(self, returncode: int, stdout: str | bytes):
+class SSHResult(BaseCmdResult[ResultOutputT]):
+    def __init__(self, returncode: int, stdout: ResultOutputT):
         super(SSHResult, self).__init__(returncode, stdout)
 
-class LocalCommandResult(BaseCmdResult):
-    def __init__(self, returncode: int, stdout: str | bytes):
+class LocalCommandResult(BaseCmdResult[ResultOutputT]):
+    def __init__(self, returncode: int, stdout: ResultOutputT):
         super(LocalCommandResult, self).__init__(returncode, stdout)
 
 def _ellide_log_lines(log: str) -> str:
@@ -77,7 +79,7 @@ def _ssh(
     decode: bool,
     options: list[str],
     multiplexing: bool,
-) -> SSHResult | SSHCommandFailed | str | bytes | None:
+) -> SSHResult[str] | SSHResult[bytes] | SSHCommandFailed | str | bytes | None:
     opts = list(options)
     opts += ['-o', 'BatchMode yes']
     opts += ['-o', 'PubkeyAcceptedKeyTypes +ssh-rsa']
@@ -137,22 +139,26 @@ def _ssh(
     if res.returncode == 255:
         return SSHCommandFailed(255, "SSH Error: %s" % output_for_errors, cmd)
 
-    output: str | bytes = res.stdout
+    output: bytes = res.stdout
     if banner_res:
         if banner_res.returncode == 255:
             return SSHCommandFailed(255, "SSH Error: %s" % banner_res.stdout.decode(errors='replace'), cmd)
         output = output[len(banner_res.stdout):]
 
-    if decode:
-        assert isinstance(output, bytes)
-        output = output.decode()
-
     if res.returncode and check:
         return SSHCommandFailed(res.returncode, output_for_errors, cmd)
 
-    if simple_output:
-        return output.strip()
-    return SSHResult(res.returncode, output)
+    if decode:
+        output_str = output.decode()
+        if simple_output:
+            return output_str.strip()
+        else:
+            return SSHResult[str](res.returncode, output_str)
+    else:
+        if simple_output:
+            return output.strip()
+        else:
+            return SSHResult[bytes](res.returncode, output)
 
 # The actual code is in _ssh().
 # This function is kept short for shorter pytest traces upon SSH failures, which are common,
@@ -173,7 +179,13 @@ def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True,
 def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True,
         simple_output: Literal[False],
         suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
-        decode: bool = True, options: List[str] = [], multiplexing: bool = True) -> SSHResult:
+        decode: Literal[True] = True, options: List[str] = [], multiplexing: bool = True) -> SSHResult[str]:
+    ...
+@overload
+def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True,
+        simple_output: Literal[False],
+        suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+        decode: Literal[False], options: List[str] = [], multiplexing: bool = True) -> SSHResult[bytes]:
     ...
 @overload
 def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True,
@@ -186,22 +198,34 @@ def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True,
         simple_output: bool = True,
         suppress_fingerprint_warnings: bool = True, background: bool = False,
         decode: bool = True, options: List[str] = [], multiplexing: bool = True) \
-        -> str | bytes | SSHResult | None:
+        -> str | bytes | SSHResult[str] | SSHResult[bytes] | None:
     ...
 def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True, simple_output: bool = True,
         suppress_fingerprint_warnings: bool = True,
         background: bool = False, decode: bool = True, options: List[str] = [], multiplexing: bool = True) \
-        -> str | bytes | SSHResult | None:
+        -> str | bytes | SSHResult[str] | SSHResult[bytes] | None:
     result_or_exc = _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
                          background, decode, options, multiplexing)
     if isinstance(result_or_exc, SSHCommandFailed):
         raise result_or_exc
     else:
         return result_or_exc
 
-def ssh_with_result(hostname_or_ip: HostAddress, cmd: str, suppress_fingerprint_warnings: bool = True,
+@overload
+def ssh_with_result(hostname_or_ip: HostAddress, cmd: str, *, decode: Literal[True] = True,
+                    suppress_fingerprint_warnings: bool = True,
+                    background: bool = False, options: List[str] = [],
+                    multiplexing: bool = True) -> SSHResult[str]:
+    ...
+@overload
+def ssh_with_result(hostname_or_ip: HostAddress, cmd: str, *, decode: Literal[False],
+                    suppress_fingerprint_warnings: bool = True,
+                    background: bool = False, options: List[str] = [],
+                    multiplexing: bool = True) -> SSHResult[bytes]:
+    ...
+def ssh_with_result(hostname_or_ip: HostAddress, cmd: str, *, suppress_fingerprint_warnings: bool = True,
                     background: bool = False, decode: bool = True, options: List[str] = [],
-                    multiplexing: bool = True) -> SSHResult:
+                    multiplexing: bool = True) -> SSHResult[str] | SSHResult[bytes]:
     result_or_exc = _ssh(hostname_or_ip, cmd, False, False, suppress_fingerprint_warnings,
                          background, decode, options, multiplexing)
     if isinstance(result_or_exc, SSHCommandFailed):
@@ -267,7 +291,15 @@ def sftp(
 
     return res
 
-def local_cmd(cmd: List[str], check: bool = True, decode: bool = True) -> LocalCommandResult:
+@overload
+def local_cmd(cmd: List[str], *, check: bool = True, decode: Literal[True] = True) -> LocalCommandResult[str]:
+    ...
+@overload
+def local_cmd(cmd: List[str], *, check: bool = True, decode: Literal[False]) -> LocalCommandResult[bytes]:
+    ...
+def local_cmd(
+    cmd: List[str], *, check: bool = True, decode: bool = True
+) -> LocalCommandResult[str] | LocalCommandResult[bytes]:
     """ Run a command locally on tester end. """
     logging.debug("[local] %s", (cmd,))
     res = subprocess.run(
@@ -280,18 +312,17 @@ def local_cmd(cmd: List[str], check: bool = True, decode: bool = True) -> LocalC
     # get a decoded version of the output in any case, replacing potential errors
     output_for_logs = res.stdout.decode(errors='replace').strip()
 
-    output: str | bytes = res.stdout
-    if decode:
-        output = res.stdout.decode()
-
     errorcode_msg = "" if res.returncode == 0 else " - Got error code: %s" % res.returncode
     command = " ".join(cmd)
     logging.debug(f"[local] {command}{errorcode_msg}{_ellide_log_lines(output_for_logs)}")
 
     if res.returncode and check:
         raise LocalCommandFailed(res.returncode, output_for_logs, command)
 
-    return LocalCommandResult(res.returncode, output)
+    if decode:
+        return LocalCommandResult[str](res.returncode, res.stdout.decode())
+    else:
+        return LocalCommandResult[bytes](res.returncode, res.stdout)
 
 def encode_powershell_command(cmd: str) -> str:
     return base64.b64encode(cmd.encode("utf-16-le")).decode("ascii")

lib/host.py

@@ -31,7 +31,7 @@
 from lib.vm import VM
 from lib.xo import xo_cli, xo_object_exists
 
-from typing import TYPE_CHECKING, Literal, TypedDict, cast, overload
+from typing import TYPE_CHECKING, Literal, TypedDict, overload
 
 if TYPE_CHECKING:
     from lib.pool import Pool
@@ -107,7 +107,13 @@ def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[True] = Tr
     @overload
     def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[False],
             suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
-            decode: bool = True, multiplexing: bool = True) -> commands.SSHResult:
+            decode: Literal[True] = True, multiplexing: bool = True) -> commands.SSHResult[str]:
+        ...
+
+    @overload
+    def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[False],
+            suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+            decode: Literal[False], multiplexing: bool = True) -> commands.SSHResult[bytes]:
         ...
 
     @overload
@@ -118,19 +124,18 @@ def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True,
 
     @overload
     def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True,
-            suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True,
-            multiplexing: bool = True) \
-            -> str | bytes | commands.SSHResult | None:
+            suppress_fingerprint_warnings: bool = True, background: Literal[False] = False,
+            decode: Literal[True] = True, multiplexing: bool = True) -> str | commands.SSHResult[str]:
         ...
 
     def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True,
             suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True,
-            multiplexing: bool = True) -> str | bytes | commands.SSHResult | None:
+            multiplexing: bool = True) -> str | bytes | commands.SSHResult[str] | commands.SSHResult[bytes] | None:
         return commands.ssh(self.hostname_or_ip, cmd, check=check, simple_output=simple_output,
                             suppress_fingerprint_warnings=suppress_fingerprint_warnings,
                             background=background, decode=decode, multiplexing=multiplexing)
 
-    def ssh_with_result(self, cmd: str) -> commands.SSHResult:
+    def ssh_with_result(self, cmd: str) -> commands.SSHResult[str]:
         # doesn't raise if the command's return is nonzero, unless there's a SSH error
         return commands.ssh_with_result(self.hostname_or_ip, cmd)
 
@@ -148,12 +153,12 @@ def xe(self, action: str, args: dict[str, str | bool | dict[str, str]] = {}, *,
 
     @overload
     def xe(self, action: str, args: dict[str, str | bool | dict[str, str]] = {}, *, check: bool = ...,
-           simple_output: Literal[False], minimal: bool = ..., force: bool = ...) -> commands.SSHResult:
+           simple_output: Literal[False], minimal: bool = ..., force: bool = ...) -> commands.SSHResult[str]:
         ...
 
     def xe(self, action: str, args: dict[str, str | bool | dict[str, str]] = {}, *, check: bool = True,
            simple_output: bool = True, minimal: bool = False, force: bool = False) \
-            -> str | commands.SSHResult:
+            -> str | commands.SSHResult[str]:
         maybe_param_minimal = '--minimal' if minimal else ''
         maybe_param_force = '--force' if force else ''
 
@@ -169,14 +174,10 @@ def stringify(key: str, value: str | bool | dict[str, str]) -> str:
 
         command: str = f'xe {action} {maybe_param_minimal} {maybe_param_force} ' + \
                        ' '.join(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))
-
-        return result
+        if simple_output:
+            return self.ssh(command, check=check, simple_output=True)
+        else:
+            return self.ssh(command, check=check, simple_output=False)
 
     @overload
     def param_get(self, param_name: str, key: str | None = ...,
@@ -239,11 +240,11 @@ def execute_script(self, script_contents: str, *, shebang: str = ..., simple_out
     @overload
     def execute_script(
         self, script_contents: str, *, shebang: str = ..., simple_output: Literal[False]
-    ) -> commands.SSHResult:
+    ) -> commands.SSHResult[str]:
         ...
 
     def execute_script(self, script_contents: str, shebang: str = 'sh',
-                       simple_output: bool = True) -> str | commands.SSHResult:
+                       simple_output: bool = True) -> str | commands.SSHResult[str]:
         with tempfile.NamedTemporaryFile('w') as script:
             os.chmod(script.name, 0o775)
             script.write('#!/usr/bin/env ' + shebang + '\n')
@@ -259,7 +260,7 @@ def execute_script(self, script_contents: str, shebang: str = 'sh',
 
             try:
                 logging.debug(f"[{self}] # Will execute this temporary script:\n{script_contents.strip()}")
-                return cast(str | commands.SSHResult, self.ssh(remote_path, simple_output=simple_output))
+                return self.ssh(remote_path, simple_output=simple_output)
             finally:
                 self.ssh(f'rm -f {remote_path}')
 

lib/vm.py

@@ -27,7 +27,7 @@
 from lib.vdi import VDI
 from lib.vif import VIF
 
-from typing import TYPE_CHECKING, Iterable, List, Literal, cast, overload
+from typing import TYPE_CHECKING, Iterable, List, Literal, overload
 
 if TYPE_CHECKING:
     from lib.host import Host
@@ -116,7 +116,12 @@ def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[True] = Tr
 
     @overload
     def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[False],
-            background: Literal[False] = False, decode: bool = True) -> commands.SSHResult:
+            background: Literal[False] = False, decode: Literal[True] = True) -> commands.SSHResult[str]:
+        ...
+
+    @overload
+    def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[False],
+            background: Literal[False] = False, decode: Literal[False]) -> commands.SSHResult[bytes]:
         ...
 
     @overload
@@ -125,18 +130,18 @@ def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True,
         ...
 
     @overload
-    def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, background: bool = False,
-            decode: bool = True) -> str | bytes | commands.SSHResult | None:
+    def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True,
+            background: Literal[False] = False, decode: Literal[True] = True) -> str | commands.SSHResult[str]:
         ...
 
     def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, background: bool = False,
-            decode: bool = True) -> str | bytes | commands.SSHResult | None:
+            decode: bool = True) -> str | bytes | commands.SSHResult[str] | commands.SSHResult[bytes] | None:
         # raises by default for any nonzero return code
         assert self.ip is not None
         return commands.ssh(self.ip, cmd, check=check, simple_output=simple_output, background=background,
                             decode=decode)
 
-    def ssh_with_result(self, cmd: str) -> commands.SSHResult:
+    def ssh_with_result(self, cmd: str) -> commands.SSHResult[str]:
         # doesn't raise if the command's return is nonzero, unless there's a SSH error
         assert self.ip is not None
         return commands.ssh_with_result(self.ip, cmd)
@@ -446,10 +451,10 @@ def execute_script(self, script_contents: str, *, simple_output: Literal[True] =
         ...
 
     @overload
-    def execute_script(self, script_contents: str, *, simple_output: Literal[False]) -> commands.SSHResult:
+    def execute_script(self, script_contents: str, *, simple_output: Literal[False]) -> commands.SSHResult[str]:
         ...
 
-    def execute_script(self, script_contents: str, simple_output: bool = True) -> str | commands.SSHResult:
+    def execute_script(self, script_contents: str, simple_output: bool = True) -> str | commands.SSHResult[str]:
         with tempfile.NamedTemporaryFile('w') as f:
             f.write(script_contents)
             f.flush()
@@ -458,7 +463,7 @@ def execute_script(self, script_contents: str, simple_output: bool = True) -> st
                 logging.debug(f"[{self.ip}] # Will execute this temporary script:\n{script_contents.strip()}")
                 # Use bash to run the script, to avoid being hit by differences between shells, for example on FreeBSD
                 # It is a documented requirement that bash is present on all test VMs.
-                res = cast(str | commands.SSHResult, self.ssh(f'bash {f.name}', simple_output=simple_output))
+                res = self.ssh(f'bash {f.name}', simple_output=simple_output)
                 return res
             finally:
                 self.ssh(f'rm -f {f.name}')
@@ -777,26 +782,29 @@ def execute_powershell_script(self, script_contents: 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,
+        simple_output: Literal[False],
+        prepend: str = "$ProgressPreference = 'SilentlyContinue';",
+    ) -> commands.SSHResult[str]:
         ...
 
     def execute_powershell_script(
             self,
             script_contents: str,
             simple_output: bool = True,
-            prepend: str = "$ProgressPreference = 'SilentlyContinue';") -> str | commands.SSHResult:
+            prepend: str = "$ProgressPreference = 'SilentlyContinue';") -> str | commands.SSHResult[str]:
         # ProgressPreference is needed to suppress any clixml progress output,
         # as it's not filtered away from stdout by default, and we're grabbing stdout.
         assert self.is_windows
         if prepend is not None:
             script_contents = prepend + script_contents
         cmd = commands.encode_powershell_command(script_contents)
-        return cast(str | commands.SSHResult, self.ssh(
+        return self.ssh(
             f"powershell.exe -nologo -noprofile -noninteractive -encodedcommand {cmd}",
             simple_output=simple_output,
-        ))
+        )
 
     def run_powershell_command(self, program: str, args: str) -> int:
         """

tests/misc/test_access_links.py

@@ -37,7 +37,6 @@ def test_access_links(host: Host, command_id: str, url_id: str) -> None:
     # Verify the download worked by comparing with local download
     # This ensures the content is accessible and identical from both locations
     local_result = commands.local_cmd(COMMAND)
-    assert isinstance(local_result.stdout, str)
 
     assert local_result.returncode == 0, f"Failed to fetch URL locally: {local_result.stdout}"