type-check devices.py exc.py features.py

Author: anon

qubesadmin/devices.py

@@ -31,8 +31,10 @@ class is implemented by an extension.
 Devices are identified by pair of (backend domain, `port_id`), where `port_id`
 is :py:class:`str`.
 """
+from __future__ import annotations
 import itertools
-from typing import Iterable
+from typing import TYPE_CHECKING
+from collections.abc import Iterable, Iterator
 
 import qubesadmin.exc
 from qubesadmin.device_protocol import (
@@ -43,6 +45,8 @@ class is implemented by an extension.
     VirtualDevice,
     AssignmentMode, DeviceInterface,
 )
+if TYPE_CHECKING:
+    from qubesadmin.vm import QubesVM
 
 
 class DeviceCollection:
@@ -55,7 +59,7 @@ class DeviceCollection:
 
     """
 
-    def __init__(self, vm, class_):
+    def __init__(self, vm: QubesVM, class_: str):
         self._vm = vm
         self._class = class_
         self._dev_cache = {}
@@ -268,7 +272,7 @@ def get_exposed_devices(self) -> Iterable[DeviceInfo]:
 
     def update_assignment(
             self, device: VirtualDevice, required: AssignmentMode
-    ):
+    ) -> None:
         """
         Update assignment of already attached device.
 
@@ -288,15 +292,15 @@ def update_assignment(
 
     __iter__ = get_exposed_devices
 
-    def clear_cache(self):
+    def clear_cache(self) -> None:
         """
         Clear cache of available devices.
         """
         self._dev_cache.clear()
         self._assignment_cache = None
         self._attachment_cache = None
 
-    def __getitem__(self, item):
+    def __getitem__(self, item: object) -> DeviceInfo:
         """Get device object with given port_id.
 
         :returns: py:class:`DeviceInfo`
@@ -316,6 +320,8 @@ def __getitem__(self, item):
                 return dev
         # if still nothing, return UnknownDevice instance for the reason
         # explained in docstring, but don't cache it
+        if not isinstance(item, str | None):
+            raise NotImplementedError
         return UnknownDevice(Port(self._vm, item, devclass=self._class))
 
 
@@ -325,21 +331,22 @@ class DeviceManager(dict):
     :param vm: VM for which we manage devices
     """
 
-    def __init__(self, vm):
+    def __init__(self, vm: QubesVM):
         super().__init__()
         self._vm = vm
 
-    def __missing__(self, key):
+    def __missing__(self, key: str) -> DeviceCollection:
         self[key] = DeviceCollection(self._vm, key)
         return self[key]
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[str]:
         return iter(self._vm.app.list_deviceclass())
 
-    def keys(self):
+
+    def keys(self) -> list[str]: # type: ignore[override]
         return self._vm.app.list_deviceclass()
 
-    def deny(self, *interfaces: Iterable[DeviceInterface]):
+    def deny(self, *interfaces: Iterable[DeviceInterface]) -> None:
         """
         Deny a device with any of the given interfaces from attaching to the VM.
         """
@@ -350,7 +357,7 @@ def deny(self, *interfaces: Iterable[DeviceInterface]):
             "".join(repr(ifc) for ifc in interfaces).encode('ascii'),
         )
 
-    def allow(self, *interfaces: Iterable[DeviceInterface]):
+    def allow(self, *interfaces: Iterable[DeviceInterface]) -> None:
         """
         Remove given interfaces from denied list.
         """
@@ -361,7 +368,7 @@ def allow(self, *interfaces: Iterable[DeviceInterface]):
             "".join(repr(ifc) for ifc in interfaces).encode('ascii'),
         )
 
-    def clear_cache(self):
+    def clear_cache(self) -> None:
         """Clear cache of all available device classes"""
         for devclass in self.values():
             devclass.clear_cache()

qubesadmin/exc.py

@@ -26,7 +26,7 @@
 class QubesException(Exception):
     """Exception that can be shown to the user"""
 
-    def __init__(self, message_format, *args, **kwargs):
+    def __init__(self, message_format: str, *args, **kwargs):
         # TODO: handle translations
         super().__init__(
             message_format % tuple(int(d) if d.isdigit() else d for d in args),
@@ -37,7 +37,7 @@ def __init__(self, message_format, *args, **kwargs):
 class QubesVMNotFoundError(QubesException, KeyError):
     """Domain cannot be found in the system"""
 
-    def __str__(self):
+    def __str__(self) -> str:
         # KeyError overrides __str__ method
         return QubesException.__str__(self)
 
@@ -139,23 +139,23 @@ class QubesMemoryError(QubesVMError, MemoryError):
 class QubesFeatureNotFoundError(QubesException, KeyError):
     """Feature not set for a given domain"""
 
-    def __str__(self):
+    def __str__(self) -> str:
         # KeyError overrides __str__ method
         return QubesException.__str__(self)
 
 
 class QubesTagNotFoundError(QubesException, KeyError):
     """Tag not set for a given domain"""
 
-    def __str__(self):
+    def __str__(self) -> str:
         # KeyError overrides __str__ method
         return QubesException.__str__(self)
 
 
 class QubesLabelNotFoundError(QubesException, KeyError):
     """Label does not exists"""
 
-    def __str__(self):
+    def __str__(self) -> str:
         # KeyError overrides __str__ method
         return QubesException.__str__(self)
 
@@ -213,7 +213,7 @@ class QubesDaemonCommunicationError(QubesException):
 class BackupRestoreError(QubesException):
     """Restoring a backup failed"""
 
-    def __init__(self, msg, backup_log=None):
+    def __init__(self, msg: str, backup_log: bytes | None=None):
         super().__init__(msg)
         self.backup_log = backup_log
 
@@ -228,7 +228,7 @@ class QubesPropertyAccessError(QubesDaemonAccessError, AttributeError):
     """Failed to read/write property value, cause is unknown (insufficient
     permissions, no such property, invalid value, other)"""
 
-    def __init__(self, prop):
+    def __init__(self, prop: str):
         super().__init__("Failed to access '%s' property" % prop)
 
 

qubesadmin/features.py

@@ -19,7 +19,16 @@
 # with this program; if not, see <http://www.gnu.org/licenses/>.
 
 '''VM features interface'''
+from __future__ import annotations
 
+import typing
+from typing import TypeVar
+from collections.abc import Iterator, Generator
+
+if typing.TYPE_CHECKING:
+    from qubesadmin.vm import QubesVM
+
+T = TypeVar('T')
 
 class Features:
     '''Manager of the features.
@@ -33,14 +42,14 @@ class Features:
     false in Python) will result in string `'0'`, which is considered true.
     '''
 
-    def __init__(self, vm):
+    def __init__(self, vm: QubesVM):
         super().__init__()
         self.vm = vm
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: str) -> None:
         self.vm.qubesd_call(self.vm.name, 'admin.vm.feature.Remove', key)
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value: object) -> None:
         if isinstance(value, bool):
             # False value needs to be serialized as empty string
             self.vm.qubesd_call(self.vm.name, 'admin.vm.feature.Set', key,
@@ -49,25 +58,30 @@ def __setitem__(self, key, value):
             self.vm.qubesd_call(self.vm.name, 'admin.vm.feature.Set', key,
                 str(value).encode())
 
-    def __getitem__(self, item):
+    def __getitem__(self, item: str) -> str:
         return self.vm.qubesd_call(
             self.vm.name, 'admin.vm.feature.Get', item).decode('utf-8')
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[str]:
         qubesd_response = self.vm.qubesd_call(self.vm.name,
             'admin.vm.feature.List')
         return iter(qubesd_response.decode('utf-8').splitlines())
 
     keys = __iter__
 
-    def items(self):
+    def items(self) -> Generator[tuple[str, str]]:
         '''Return iterable of pairs (feature, value)'''
         for key in self:
             yield key, self[key]
 
     NO_DEFAULT = object()
 
-    def get(self, item, default=None):
+    @typing.overload
+    def get(self, item: str) -> str | None: ...
+    @typing.overload
+    def get(self, item: str, default: T) -> str | T: ...
+    # Overloaded to handle default None return type
+    def get(self, item: str, default: object = None) -> object:
         '''Get a feature, return default value if missing.'''
         try:
             return self[item]
@@ -76,7 +90,13 @@ def get(self, item, default=None):
                 raise
             return default
 
-    def check_with_template(self, feature, default=None):
+    @typing.overload
+    def check_with_template(self, item: str) -> str | None: ...
+    @typing.overload
+    def check_with_template(self, item: str, default: T) -> str | T: ...
+    # Overloaded to handle default None return type
+    def check_with_template(self, feature: str,
+                            default: object = None) -> object:
         ''' Check if the vm's template has the specified feature. '''
         try:
             qubesd_response = self.vm.qubesd_call(