WIP: address mypy issues

pyproject.toml

@@ -8,7 +8,8 @@ dependencies = [
   "snakemake-interface-common>=1.12.0",
   "wrapt>=1.15.0",
   "reretry>=0.11.8",
-  "throttler>=1.2.2"
+  "throttler>=1.2.2",
+  "humanfriendly"
 ]
 
 [[project.authors]]

snakemake_interface_storage_plugins/io.py

@@ -6,6 +6,7 @@
 from abc import abstractmethod
 import re
 from typing import Dict
+from typing import Optional
 
 
 WILDCARD_REGEX = re.compile(
@@ -26,7 +27,7 @@
 )
 
 
-def get_constant_prefix(pattern: str, strip_incomplete_parts: bool = False):
+def get_constant_prefix(pattern: str, strip_incomplete_parts: bool = False) -> str:
     """Return constant prefix of a pattern, removing everything from the first
     wildcard on.
 
@@ -53,23 +54,31 @@ def get_constant_prefix(pattern: str, strip_incomplete_parts: bool = False):
 
 class Mtime:
     __slots__ = ["_local", "_local_target", "_storage"]
+    _local: Optional[float]
+    _local_target: Optional[float]
+    _storage: Optional[float]
 
-    def __init__(self, local=None, local_target=None, storage=None):
+    def __init__(
+        self,
+        local: Optional[float] = None,
+        local_target: Optional[float] = None,
+        storage: Optional[float] = None,
+    ):
         self._local = local
         self._local_target = local_target
         self._storage = storage
 
-    def local_or_storage(self, follow_symlinks=False):
+    def local_or_storage(self, follow_symlinks: bool = False) -> Optional[float]:
         if self._storage is not None:
             return self._storage
         return self.local(follow_symlinks=follow_symlinks)
 
     def storage(
         self,
-    ):
+    ) -> Optional[float]:
         return self._storage
 
-    def local(self, follow_symlinks=False):
+    def local(self, follow_symlinks: bool = False) -> Optional[float]:
         if follow_symlinks and self._local_target is not None:
             return self._local_target
         return self._local

snakemake_interface_storage_plugins/registry/__init__.py

@@ -20,14 +20,14 @@
 from snakemake_interface_storage_plugins.storage_provider import StorageProviderBase
 
 
-class StoragePluginRegistry(PluginRegistryBase):
+class StoragePluginRegistry(PluginRegistryBase[Plugin]):
     """This class is a singleton that holds all registered executor plugins."""
 
     def get_registered_read_write_plugins(self) -> List[str]:
         return [
             plugin.name
             for plugin in self.plugins.values()
-            if plugin.storage_provider.is_read_write
+            if plugin.storage_provider.is_read_write is True
         ]
 
     @property

snakemake_interface_storage_plugins/registry/plugin.py

@@ -4,7 +4,7 @@
 __license__ = "MIT"
 
 from dataclasses import dataclass
-from typing import Optional, Type
+from typing import Optional, Type, TYPE_CHECKING
 from snakemake_interface_storage_plugins.settings import (
     StorageProviderSettingsBase,
 )
@@ -18,10 +18,15 @@
 )
 
 
+if TYPE_CHECKING:
+    from snakemake_interface_storage_plugins.storage_provider import StorageProviderBase
+    from snakemake_interface_storage_plugins.storage_object import StorageObjectBase
+
+
 @dataclass
-class Plugin(PluginBase):
-    storage_provider: object
-    storage_object: object
+class Plugin(PluginBase[StorageProviderSettingsBase]):
+    storage_provider: Type["StorageProviderBase"]
+    storage_object: Type["StorageObjectBase"]
     _storage_settings_cls: Optional[Type[StorageProviderSettingsBase]]
     _name: str
 
@@ -30,18 +35,19 @@ def support_tagged_values(self) -> bool:
         return True
 
     @property
-    def name(self):
+    def name(self) -> str:
         return self._name
 
     @property
-    def cli_prefix(self):
+    def cli_prefix(self) -> str:
         return "storage-" + self.name.replace(common.storage_plugin_module_prefix, "")
 
     @property
-    def settings_cls(self):
+    def settings_cls(self) -> Optional[Type[StorageProviderSettingsBase]]:
         return self._storage_settings_cls
 
-    def is_read_write(self):
+    @property
+    def is_read_write(self) -> bool:
         return issubclass(self.storage_object, StorageObjectWrite) and issubclass(
             self.storage_object, StorageObjectRead
         )

snakemake_interface_storage_plugins/storage_object.py

@@ -8,7 +8,7 @@
 from abc import ABC, abstractmethod
 from pathlib import Path
 import shutil
-from typing import Iterable, Optional
+from typing import Iterable, Optional, AsyncContextManager, Dict, TypeVar, Generic
 
 from wrapt import ObjectProxy
 from reretry import retry
@@ -34,56 +34,64 @@ class StaticStorageObjectProxy(ObjectProxy):
 
     """
 
-    def exists(self):
+    def exists(self) -> bool:
         return True
 
     def mtime(self) -> float:
         return float("-inf")
 
-    def is_newer(self, time):
+    def is_newer(self, time: float) -> bool:
         return False
 
-    def __copy__(self):
+    def __copy__(self) -> "StaticStorageObjectProxy":
         copied_wrapped = copy.copy(self.__wrapped__)
         return type(self)(copied_wrapped)
 
-    def __deepcopy__(self, memo):
+    def __deepcopy__(self, memo: Dict) -> "StaticStorageObjectProxy":
         copied_wrapped = copy.deepcopy(self.__wrapped__, memo)
         return type(self)(copied_wrapped)
 
 
-class StorageObjectBase(ABC):
+TStorageProviderBase = TypeVar("TStorageProviderBase", bound=StorageProviderBase)
+
+
+class StorageObjectBase(ABC, Generic[TStorageProviderBase]):
     """This is an abstract class to be used to derive storage object classes for
     different cloud storage providers. For example, there could be classes for
     interacting with Amazon AWS S3 and Google Cloud Storage, both derived from this
     common base class.
     """
+    query: str
+    keep_local: bool
+    retrieve: bool
+    provider: TStorageProviderBase
+    print_query: str
+    _overwrite_local_path: Optional[Path] = None
+    _is_ondemand_eligible: bool = False
 
     def __init__(
         self,
         query: str,
         keep_local: bool,
         retrieve: bool,
-        provider: StorageProviderBase,
-    ):
-        self.query: str = query
-        self.keep_local: bool = keep_local
-        self.retrieve: bool = retrieve
-        self.provider: StorageProviderBase = provider
-        self.print_query: str = self.provider.safe_print(self.query)
-        self._overwrite_local_path: Optional[Path] = None
-        self._is_ondemand_eligible: bool = False
+        provider: TStorageProviderBase,
+    ) -> None:
+        self.query = query
+        self.keep_local  = keep_local
+        self.retrieve  = retrieve
+        self.provider = provider
+        self.print_query = self.provider.safe_print(self.query)
         self.__post_init__()
 
-    def __post_init__(self):  # noqa B027
+    def __post_init__(self) -> None:  # noqa B027
         pass
 
     @property
     def is_ondemand_eligible(self) -> bool:
         return self._is_ondemand_eligible and not self.keep_local
 
     @is_ondemand_eligible.setter
-    def is_ondemand_eligible(self, value: bool):
+    def is_ondemand_eligible(self, value: bool) -> None:
         self._is_ondemand_eligible = value
 
     def set_local_path(self, path: Path) -> None:
@@ -92,7 +100,7 @@ def set_local_path(self, path: Path) -> None:
 
     def is_valid_query(self) -> bool:
         """Return True is the query is valid for this storage provider."""
-        return self.provider.is_valid_query(self.query)
+        return bool(self.provider.is_valid_query(self.query))
 
     def local_path(self) -> Path:
         """Return the local path that would represent the query."""
@@ -116,13 +124,13 @@ def local_suffix(self) -> str:
         # part and any optional parameters if that does not hamper the uniqueness.
         ...
 
-    def _rate_limiter(self, operation: Operation):
+    def _rate_limiter(self, operation: Operation) -> AsyncContextManager:
         return self.provider.rate_limiter(self.query, operation)
 
 
 class StorageObjectRead(StorageObjectBase):
     @abstractmethod
-    async def inventory(self, cache: IOCacheStorageInterface):
+    async def inventory(self, cache: IOCacheStorageInterface) -> None:
         """From this file, try to find as much existence and modification date
         information as possible.
         """
@@ -134,7 +142,7 @@ async def inventory(self, cache: IOCacheStorageInterface):
     def get_inventory_parent(self) -> Optional[str]: ...
 
     @abstractmethod
-    def cleanup(self):
+    def cleanup(self) -> None:
         """Perform local cleanup of any remainders of the storage object."""
         ...
 
@@ -157,7 +165,7 @@ def local_footprint(self) -> int:
         return self.size()
 
     @abstractmethod
-    def retrieve_object(self):
+    def retrieve_object(self) -> None:
         """Ensure that the object is accessible locally under self.local_path()
 
         Optionally, this can make use of the attribute self.is_ondemand_eligible,
@@ -194,7 +202,7 @@ async def managed_exists(self) -> bool:
         except Exception as e:
             raise WorkflowError(f"Failed to check existence of {self.print_query}", e)
 
-    async def managed_retrieve(self):
+    async def managed_retrieve(self) -> None:
         await self.wait_for_free_space()
         try:
             self.local_path().parent.mkdir(parents=True, exist_ok=True)
@@ -223,7 +231,7 @@ async def managed_local_footprint(self) -> int:
                 e,
             )
 
-    async def wait_for_free_space(self):
+    async def wait_for_free_space(self) -> None:
         """Wait for free space on the disk."""
         size = await self.managed_local_footprint()
         disk_free = get_disk_free(self.local_path())
@@ -253,19 +261,19 @@ async def wait_for_free_space(self):
             raise WorkflowError(
                 f"Cannot store {self.local_path()} "
                 f"({format_size(size)} > {format_size(disk_free)}), "
-                f"waited {format_timespan(self.provider.wait_for_free_local_storage)} "
+                f"waited {format_timespan(self.provider.wait_for_free_local_storage or 0)} "
                 "for more space."
             )
 
 
 class StorageObjectWrite(StorageObjectBase):
     @abstractmethod
-    def store_object(self): ...
+    def store_object(self) -> None: ...
 
     @abstractmethod
-    def remove(self): ...
+    def remove(self) -> None: ...
 
-    async def managed_remove(self):
+    async def managed_remove(self) -> None:
         try:
             async with self._rate_limiter(Operation.REMOVE):
                 self.remove()
@@ -274,7 +282,7 @@ async def managed_remove(self):
                 f"Failed to remove storage object {self.print_query}", e
             )
 
-    async def managed_store(self):
+    async def managed_store(self) -> None:
         try:
             async with self._rate_limiter(Operation.STORE):
                 self.store_object()
@@ -295,11 +303,11 @@ def list_candidate_matches(self) -> Iterable[str]:
 
 class StorageObjectTouch(StorageObjectBase):
     @abstractmethod
-    def touch(self):
+    def touch(self) -> None:
         """Touch the object."""
         ...
 
-    async def managed_touch(self):
+    async def managed_touch(self) -> None:
         try:
             async with self._rate_limiter(Operation.TOUCH):
                 self.touch()