Merge pull request #284 from xcp-ng/add-some-type-hints

add some type hints

Author: ydirson

.github/workflows/jobs-check.yml

@@ -37,3 +37,21 @@ jobs:
       run: cp data.py-dist data.py
     - name: Check with mypy
       run: mypy lib/ tests/
+      
+  pyright:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: 3.8
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements/base.txt
+        pip install pyright
+    - name: Create a dummy data.py
+      run: cp data.py-dist data.py
+    - name: Check with pyright
+      run: pyright lib/ # tests/

data.py-dist

@@ -12,6 +12,7 @@ HOST_DEFAULT_PASSWORD = ""
 # that the tests will create or import, such as VMs and SRs.
 # Default value: [your login/user]
 # OBJECTS_NAME_PREFIX = "[TEST]"
+OBJECTS_NAME_PREFIX = None
 
 # Override settings for specific hosts
 # skip_xo_config allows to not touch XO's configuration regarding the host

lib/basevm.py

@@ -1,6 +1,10 @@
 import logging
 
+from typing import Any, Optional, TYPE_CHECKING, Union
+
 import lib.commands as commands
+if TYPE_CHECKING:
+    import lib.host
 
 from lib.common import _param_add, _param_clear, _param_get, _param_remove, _param_set
 from lib.sr import SR
@@ -10,33 +14,40 @@ class BaseVM:
 
     xe_prefix = "vm"
 
-    def __init__(self, uuid, host):
+    def __init__(self, uuid, host: 'lib.host.Host'):
         logging.info("New %s: %s", type(self).__name__, uuid)
         self.uuid = uuid
         self.host = host
 
-    def param_get(self, param_name, key=None, accept_unknown_key=False):
+    def param_get(self, param_name: str, key: Optional[str] = None,
+                  accept_unknown_key: bool = False) -> Union[str, bool, None]:
         return _param_get(self.host, self.xe_prefix, self.uuid,
                           param_name, key, accept_unknown_key)
 
-    def param_set(self, param_name, value, key=None):
+    def param_set(self, param_name: str, value: Any, key: Optional[str] = None) -> None:
         _param_set(self.host, self.xe_prefix, self.uuid,
                    param_name, value, key)
 
-    def param_remove(self, param_name, key, accept_unknown_key=False):
+    def param_remove(self, param_name: str, key: str, accept_unknown_key=False) -> None:
         _param_remove(self.host, self.xe_prefix, self.uuid,
                       param_name, key, accept_unknown_key)
 
-    def param_add(self, param_name, value, key=None):
+    def param_add(self, param_name: str, value: str, key=None) -> None:
         _param_add(self.host, self.xe_prefix, self.uuid,
                    param_name, value, key)
 
-    def param_clear(self, param_name):
+    def param_clear(self, param_name: str) -> None:
         _param_clear(self.host, self.xe_prefix, self.uuid,
                      param_name)
 
-    def name(self):
-        return self.param_get('name-label')
+    def name(self) -> str:
+        n = self.param_get('name-label')
+        assert isinstance(n, str)
+        return n
+
+    # @abstractmethod
+    def _disk_list(self):
+        raise NotImplementedError()
 
     def vdi_uuids(self, sr_uuid=None):
         output = self._disk_list()
@@ -54,23 +65,17 @@ def vdi_uuids(self, sr_uuid=None):
                 vdis_on_sr.append(vdi)
         return vdis_on_sr
 
-    def destroy_vdi(self, vdi_uuid):
+    def destroy_vdi(self, vdi_uuid: str) -> None:
         self.host.xe('vdi-destroy', {'uuid': vdi_uuid})
 
-    # FIXME: move this method and the above back to class VM if not useful in Snapshot class?
-    def destroy(self):
-        for vdi_uuid in self.vdi_uuids():
-            self.destroy_vdi(vdi_uuid)
-        self._destroy()
-
     def all_vdis_on_host(self, host):
         for vdi_uuid in self.vdi_uuids():
             sr = SR(self.host.pool.get_vdi_sr_uuid(vdi_uuid), self.host.pool)
             if not sr.attached_to_host(host):
                 return False
         return True
 
-    def all_vdis_on_sr(self, sr):
+    def all_vdis_on_sr(self, sr) -> bool:
         for vdi_uuid in self.vdi_uuids():
             if self.host.pool.get_vdi_sr_uuid(vdi_uuid) != sr.uuid:
                 return False
@@ -84,7 +89,7 @@ def get_sr(self):
         assert sr.attached_to_host(self.host)
         return sr
 
-    def export(self, filepath, compress='none'):
+    def export(self, filepath, compress='none') -> None:
         logging.info("Export VM %s to %s with compress=%s" % (self.uuid, filepath, compress))
         params = {
             'uuid': self.uuid,

lib/commands.py

@@ -2,11 +2,13 @@
 import logging
 import shlex
 import subprocess
+from typing import List, Literal, overload, Union
 
-import lib.config as config
 
+import lib.config as config
 from lib.netutil import wrap_ip
 
+
 class BaseCommandFailed(Exception):
     __slots__ = 'returncode', 'stdout', 'cmd'
 
@@ -61,7 +63,7 @@ def _ellide_log_lines(log):
     return "\n{}".format("\n".join(reduced_message))
 
 def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
-         background, decode, options):
+         background, decode, options) -> Union[SSHResult, SSHCommandFailed, str, bytes, None]:
     opts = list(options)
     opts.append('-o "BatchMode yes"')
     if suppress_fingerprint_warnings:
@@ -80,6 +82,7 @@ def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warning
     ssh_cmd = f"ssh root@{hostname_or_ip} {' '.join(opts)} {shlex.quote(command)}"
 
     # Fetch banner and remove it to avoid stdout/stderr pollution.
+    banner_res = None
     if config.ignore_ssh_banner:
         banner_res = subprocess.run(
             "ssh root@%s %s '%s'" % (hostname_or_ip, ' '.join(opts), '\n'),
@@ -97,9 +100,10 @@ def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warning
         stderr=subprocess.STDOUT
     )
     if background:
-        return True, None
+        return None
 
     stdout = []
+    assert process.stdout is not None
     for line in iter(process.stdout.readline, b''):
         readable_line = line.decode(errors='replace').strip()
         stdout.append(line)
@@ -112,34 +116,73 @@ def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warning
 
     # Even if check is False, we still raise in case of return code 255, which means a SSH error.
     if res.returncode == 255:
-        return False, SSHCommandFailed(255, "SSH Error: %s" % output_for_errors, command)
+        return SSHCommandFailed(255, "SSH Error: %s" % output_for_errors, command)
 
-    output = res.stdout
-    if config.ignore_ssh_banner:
+    output: Union[bytes, str] = res.stdout
+    if banner_res:
         if banner_res.returncode == 255:
-            return False, SSHCommandFailed(255, "SSH Error: %s" % banner_res.stdout.decode(errors='replace'), command)
+            return SSHCommandFailed(255, "SSH Error: %s" % banner_res.stdout.decode(errors='replace'), command)
         output = output[len(banner_res.stdout):]
 
     if decode:
+        assert isinstance(output, bytes)
         output = output.decode()
 
     if res.returncode and check:
-        return False, SSHCommandFailed(res.returncode, output_for_errors, command)
+        return SSHCommandFailed(res.returncode, output_for_errors, command)
 
     if simple_output:
-        return True, output.strip()
-    return True, SSHResult(res.returncode, output)
+        return output.strip()
+    return SSHResult(res.returncode, output)
 
 # 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,
+        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,
+        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,
+        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],
+        decode: bool = True, options: List[str] = []) -> None:
+    ...
+@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,
+        decode: bool = True, options: List[str] = []) \
+        -> Union[str, bytes, SSHResult, None]:
+    ...
+def ssh(hostname_or_ip, cmd, *, check=True, simple_output=True,
+        suppress_fingerprint_warnings=True,
         background=False, decode=True, options=[]):
-    success, result_or_exc = _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
-                                  background, decode, options)
-    if not success:
+    result_or_exc = _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
+                         background, decode, options)
+    if isinstance(result_or_exc, SSHCommandFailed):
+        raise result_or_exc
+    else:
+        return result_or_exc
+
+def ssh_with_result(hostname_or_ip, cmd, suppress_fingerprint_warnings=True,
+                    background=False, decode=True, options=[]) -> SSHResult:
+    result_or_exc = _ssh(hostname_or_ip, cmd, False, False, suppress_fingerprint_warnings,
+                         background, decode, options)
+    if isinstance(result_or_exc, SSHCommandFailed):
         raise result_or_exc
-    return result_or_exc
+    elif isinstance(result_or_exc, SSHResult):
+        return result_or_exc
+    assert False, "unexpected type"
 
 def scp(hostname_or_ip, src, dest, check=True, suppress_fingerprint_warnings=True, local_dest=False):
     opts = '-o "BatchMode yes"'
@@ -173,6 +216,7 @@ def scp(hostname_or_ip, src, dest, check=True, suppress_fingerprint_warnings=Tru
     return res
 
 def sftp(hostname_or_ip, cmds, check=True, suppress_fingerprint_warnings=True):
+    opts = ''
     if suppress_fingerprint_warnings:
         # Suppress warnings and questions related to host key fingerprints
         # because on a test network IPs get reused, VMs are reinstalled, etc.

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
@@ -23,10 +26,13 @@ def vm_image(vm_key):
     return url
 
 def prefix_object_name(label):
+    name_prefix = None
     try:
         from data import OBJECTS_NAME_PREFIX
         name_prefix = OBJECTS_NAME_PREFIX
     except ImportError:
+        pass
+    if name_prefix is None:
         name_prefix = f"[{getpass.getuser()}]"
     return f"{name_prefix} {label}"
 
@@ -154,9 +160,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) -> Union[str, bool, None]:
     """ 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:

lib/efi.py

@@ -218,6 +218,7 @@ def sign(payload, key_file, cert_file):
     """Returns a signed PKCS7 of payload signed by key and cert."""
     with open(key_file, 'rb') as f:
         priv_key = serialization.load_pem_private_key(f.read(), password=None)
+        assert isinstance(priv_key, (pkcs7.PKCS7PrivateKeyTypes))
 
     with open(cert_file, 'rb') as f:
         cert = x509.load_pem_x509_certificate(f.read())