Merge pull request #401 from xcp-ng/gln/misc-typehints

Author: glehmann

jobs.py

@@ -7,7 +7,18 @@
 
 from lib.commands import ssh
 
-JOBS = {
+from typing import NotRequired, TypedDict, cast
+
+class JobData(TypedDict):
+    description: str
+    requirements: list[str]
+    nb_pools: int
+    params: dict[str, str]
+    paths: list[str]
+    markers: NotRequired[str]
+    name_filter: NotRequired[str]
+
+JOBS: dict[str, JobData] = {
     "main": {
         "description": "a group of not-too-long tests that run either without a VM, or with a single small one",
         "requirements": [
@@ -472,12 +483,15 @@
     "tests/storage/zfsvol/test_zfsvol_sr.py::TestZfsvolVm::test_quicktest",
 ]
 
+VmDef = str | tuple[str, str]
+VMSDef = dict[str, dict[str, VmDef | list[VmDef]]]
+
 # Returns the vm filename or None if a host_version is passed and matches the one specified
 # with the vm filename in vm_data.py. ex: ("centos6-32-hvm-created_8.2-zstd.xva", "8\.2\..*")
-def filter_vm(vm, host_version):
+def filter_vm(vm: VmDef, host_version: str | None) -> str | None:
     import re
 
-    if type(vm) is tuple:
+    if isinstance(vm, tuple):
         if len(vm) != 2:
             print(f"ERROR: VM definition from vm_data.py is a tuple so it should contain exactly two items:\n{vm}")
             sys.exit(1)
@@ -496,34 +510,34 @@ def filter_vm(vm, host_version):
 
     return vm
 
-def get_vm_or_vms_refs(handle, host_version=None):
+def get_vm_or_vms_refs(handle: str, host_version: str | None = None) -> str | list[str]:
     try:
-        from vm_data import VMS
+        from vm_data import VMS as VMS_untyped
     except ImportError:
         print("ERROR: Could not import VMS from vm_data.py.")
-        print("Get the latest vm_data.py from XCP-ng's internal lab or copy data.py-dist and fill with your VM refs.")
+        print("Get the latest vm_data.py from XCP-ng's internal lab or copy vm_data.py-dist and fill"
+              " with your VM refs.")
         print("You may also bypass this error by providing your own --vm parameter(s).")
         sys.exit(1)
 
+    VMS = cast(VMSDef, VMS_untyped)
     category, key = handle.split("/")
-    if category not in VMS or not VMS[category].get(key):
+    if category not in VMS or key not in VMS[category]:
         print(f"ERROR: Could not find VMS['{category}']['{key}'] in vm_data.py, or it's empty.")
         print("You need to update your local vm_data.py.")
         print("You may also bypass this error by providing your own --vm parameter(s).")
         sys.exit(1)
 
-    if type(VMS[category][key]) is list:
+    vms: str | list[str] | None = []
+    vms_unfiltered = VMS[category][key]
+    if isinstance(vms_unfiltered, list):
         # Multi VMs
-        vms = list()
-        for vm in VMS[category][key]:
-            xva = filter_vm(vm, host_version)
-            if xva is not None:
-                vms.append(xva)
-        if len(vms) == 0:
+        vms = [xva for vm in vms_unfiltered if (xva := filter_vm(vm, host_version)) is not None]
+        if vms == []:
             vms = None
-    else:
+    elif isinstance(vms_unfiltered, str):
         # Single VMs
-        vms = filter_vm(VMS[category][key], host_version)
+        vms = filter_vm(vms_unfiltered, host_version)
 
     if vms is None:
         print(f"ERROR: Could not find VMS['{category}']['{key}'] for host version {host_version}.")
@@ -533,7 +547,8 @@ def get_vm_or_vms_refs(handle, host_version=None):
 
     return vms
 
-def build_pytest_cmd(job_data, hosts=None, host_version=None, pytest_args=[]):
+def build_pytest_cmd(job_data: JobData, hosts: str | None = None, host_version: str | None = None,
+                     pytest_args: list[str] = []) -> list[str]:
     markers = job_data.get("markers", None)
     name_filter = job_data.get("name_filter", None)
 
@@ -543,13 +558,12 @@ def build_pytest_cmd(job_data, hosts=None, host_version=None, pytest_args=[]):
     if hosts is not None:
         try:
             host = hosts.split(',')[0]
-            cmd = "lsb_release -sr"
-            host_version = ssh(host, cmd)
+            host_version = ssh(host, "lsb_release -sr")
         except Exception as e:
             print(e, file=sys.stderr)
 
-    def _join_pytest_args(arg, option):
-        cli_args = []
+    def _join_pytest_args(arg: str | None, option: str) -> str | None:
+        cli_args: list[str] = []
         try:
             while True:
                 i = pytest_args.index(option)
@@ -600,21 +614,21 @@ def _join_pytest_args(arg, option):
     cmd += pytest_args
     return cmd
 
-def action_list(args):
+def action_list(args: argparse.Namespace) -> None:
     for job, data in JOBS.items():
         print(f"{job}: {data['description']}")
 
-def action_show(args):
+def action_show(args: argparse.Namespace) -> None:
     print(json.dumps(JOBS[args.job], indent=4))
 
-def action_collect(args):
+def action_collect(args: argparse.Namespace) -> None:
     cmd = build_pytest_cmd(JOBS[args.job], None, args.host_version, ["--collect-only"] + args.pytest_args)
     subprocess.run(cmd)
 
-def action_check(args):
+def action_check(args: argparse.Namespace) -> None:
     error = False
 
-    def extract_tests(cmd):
+    def extract_tests(cmd: list[str]) -> set[str]:
         tests = set()
         res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         if res.returncode != 0 and res.returncode != 5: # 5 means no test found
@@ -663,6 +677,7 @@ def extract_tests(cmd):
     multi_vm_tests = extract_tests(["pytest", "--collect-only", "-q", "-m", "multi_vms"]) - broken_tests
     job_tests = set()
     for job_data in JOBS.values():
+        assert isinstance(job_data["params"], dict)
         if "--vm[]" in job_data["params"]:
             job_tests |= extract_tests(build_pytest_cmd(job_data, None, None, ["--collect-only", "-q", "--vm=a_vm"]))
     tests_missing = sorted(list(multi_vm_tests - job_tests))
@@ -676,23 +691,25 @@ def extract_tests(cmd):
     if error:
         sys.exit(1)
 
-def action_run(args):
+def action_run(args: argparse.Namespace) -> None:
     cmd = build_pytest_cmd(JOBS[args.job], args.hosts, None, args.pytest_args)
     print(subprocess.list2cmdline(cmd))
     if args.print_only:
         return
 
     # check that enough pool masters have been provided
     nb_pools = len(args.hosts.split(","))
-    if nb_pools < JOBS[args.job]["nb_pools"]:
-        print(f"Error: only {nb_pools} master host(s) provided, {JOBS[args.job]['nb_pools']} required.")
+    job_nb_pools = JOBS[args.job]["nb_pools"]
+    assert isinstance(job_nb_pools, int)
+    if nb_pools < job_nb_pools:
+        print(f"Error: only {nb_pools} master host(s) provided, {job_nb_pools} required.")
         sys.exit(1)
 
     res = subprocess.run(cmd)
     if res.returncode:
         sys.exit(1)
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser(description="Manage test jobs")
     subparsers = parser.add_subparsers(dest="action", metavar="action")
     subparsers.required = True

scripts/guests/windows/vm_compact.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from __future__ import annotations
+
 import argparse
 import logging
 import os
@@ -14,6 +16,8 @@
 from lib.vbd import VBD
 from lib.vm import VM
 
+from typing import Generator
+
 # Tool to compact Windows VMs prior to test XVA creation.
 # Requirements:
 # - ntfsclone (part of ntfs-3g) and sfdisk should be installed.
@@ -22,7 +26,7 @@
 # - The tool should be able to connect to and control the pool via XAPI.
 
 
-def get_partitions(part_table: str):
+def get_partitions(part_table: str) -> Generator[dict[str, str], None, None]:
     for line in part_table.splitlines():
         match = re.match(
             (
@@ -39,13 +43,13 @@ def get_partitions(part_table: str):
         yield match.groupdict()
 
 
-def get_device_path(vbd: VBD):
+def get_device_path(vbd: VBD) -> str:
     device = vbd.param_get("device")
     assert device
     return "/dev/" + device
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", required=True, metavar="HOST", help="host IP or hostname")
     parser.add_argument("--vm", nargs="+", required=True, metavar="VM_UUID", help="UUIDs of VMs to compact")

scripts/install_xcpng.py

@@ -1,5 +1,6 @@
-#!/usr/bin/env python3
+from __future__ import annotations
 
+#!/usr/bin/env python3
 import argparse
 import atexit
 import logging
@@ -25,7 +26,8 @@
 
 logging.basicConfig(format='[%(levelname)s] %(message)s', level=logging.INFO)
 
-def generate_answerfile(directory, installer, hostname_or_ip, target_hostname, action, hdd, netinstall_gpg_check):
+def generate_answerfile(directory: str, installer: str, hostname_or_ip: str, target_hostname: str | None, action: str,
+                        hdd: str, netinstall_gpg_check: str) -> None:
     password = host_data(hostname_or_ip)['password']
     cmd = ['openssl', 'passwd', '-6', password]
     res = subprocess.run(cmd, stdout=subprocess.PIPE)
@@ -69,33 +71,33 @@ def generate_answerfile(directory, installer, hostname_or_ip, target_hostname, a
         else:
             raise Exception(f"Unknown action: `{action}`")
 
-def is_ip_active(ip):
+def is_ip_active(ip: str) -> bool:
     return not os.system(f"ping -c 3 -W 10 {ip} > /dev/null 2>&1")
 
-def is_ssh_up(ip):
+def is_ssh_up(ip: str) -> bool:
     try:
         ssh(ip, 'true', options=['-o', 'ConnectTimeout 10'])
         return True
     except SSHCommandFailed:
         # probably not up yet
         return False
 
-def get_new_host_ip(mac_address):
+def get_new_host_ip(mac_address: str) -> str | None:
     candidate_ips = pxe.arp_addresses_for(mac_address)
     logging.debug("Candidate IPs: " + ", ".join(candidate_ips))
     for ip in candidate_ips:
         if is_ip_active(ip) and is_ssh_up(ip):
             return ip
     return None
 
-def is_new_host_ready(ip_address):
+def is_new_host_ready(ip_address: str) -> bool:
     try:
         output = ssh(ip_address, 'xe host-list enabled=true --minimal')
         return is_uuid(output)
     except Exception:
         return False
 
-def check_mac_address(host: Host, mac_address):
+def check_mac_address(host: Host, mac_address: str) -> None:
     bridge = host.inventory['MANAGEMENT_INTERFACE']
     host_mac_address = host.ssh(f'cat /sys/class/net/{bridge}/address')
     if mac_address != host_mac_address:
@@ -104,15 +106,15 @@ def check_mac_address(host: Host, mac_address):
             f"Expected: `{mac_address}`"
         )
 
-def url_checker(url):
+def url_checker(url: str) -> None:
     try:
         response = requests.get(url)
         if not response:
             raise Exception(f"{url}: URL is not reachable, status_code: {response.status_code}")
     except requests.exceptions.RequestException as e:
         raise SystemExit(f"{url}: URL is not reachable\nErr: {e}")
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "host",