fix: use Protocol type

Author: johanneskoester

snakemake_interface_software_deployment_plugins/__init__.py

@@ -9,7 +9,7 @@
 import hashlib
 from pathlib import Path
 import shutil
-from typing import Any, Dict, Iterable, Optional, Self, Tuple, Type, Union
+from typing import Any, ClassVar, Dict, Iterable, Optional, Protocol, Self, Tuple, Type, Union, runtime_checkable
 import subprocess as sp
 
 from snakemake_interface_software_deployment_plugins.settings import (
@@ -123,8 +123,19 @@ def __str__(self) -> str:
         ...
 
 
-class EnvBase(ABC):
-    _cache: Dict[Tuple[Type["EnvBase"], Optional["EnvBase"]], Any] = {}
+class EnvBase(Protocol):
+    _cache: ClassVar[Dict[Tuple[Type["EnvBase"], Optional["EnvBase"]], Any]] = {}
+    spec: EnvSpecBase
+    within: Optional["EnvBase"]
+    settings: Optional[SoftwareDeploymentSettingsBase]
+    shell_executable: str
+    tempdir: Path
+    _cache_prefix: Path
+    _deployment_prefix: Path
+    _pinfile_prefix: Path
+    _managed_hash_store: Optional[str] = None
+    _managed_deployment_hash_store: Optional[str] = None
+    _obj_hash: Optional[int] = None
 
     def __init__(
         self,
@@ -142,9 +153,6 @@ def __init__(
         self.settings: Optional[SoftwareDeploymentSettingsBase] = settings
         self.shell_executable: str = shell_executable
         self.tempdir = tempdir
-        self._managed_hash_store: Optional[str] = None
-        self._managed_deployment_hash_store: Optional[str] = None
-        self._obj_hash: Optional[int] = None
         self._deployment_prefix: Path = deployment_prefix
         self._cache_prefix: Path = cache_prefix
         self._pinfile_prefix: Path = pinfile_prefix
@@ -177,6 +185,21 @@ def decorate_shellcmd(self, cmd: str) -> str:
         """
         ...
 
+    def is_deployable(self) -> bool:
+        """Overwrite this in case the deployability of the environment depends on
+        the spec or settings."""
+        return isinstance(self, DeployableEnvBase)
+
+    def is_pinnable(self) -> bool:
+        """Overwrite this in case the pinability of the environment depends on
+        the spec or settings."""
+        return isinstance(self, PinnableEnvBase)
+
+    def is_cacheable(self) -> bool:
+        """Overwrite this in case the cacheability of the environment depends on
+        the spec or settings."""
+        return isinstance(self, CacheableEnvBase)
+
     @abstractmethod
     def record_hash(self, hash_object) -> None:
         """Update given hash object (using hash_object.update()) such that it changes
@@ -241,7 +264,8 @@ def __eq__(self, other) -> bool:
         )
 
 
-class PinnableEnvBase(ABC):
+@runtime_checkable
+class PinnableEnvBase(EnvBase, Protocol):
     @classmethod
     @abstractmethod
     def pinfile_extension(cls) -> str: ...
@@ -256,7 +280,6 @@ async def pin(self) -> None:
 
     @property
     def pinfile(self) -> Path:
-        assert isinstance(self, EnvBase)
         ext = self.pinfile_extension()
         if not ext.startswith("."):
             raise ValueError("pinfile_extension must start with a dot.")
@@ -265,7 +288,8 @@ def pinfile(self) -> Path:
         )
 
 
-class CacheableEnvBase(ABC):
+@runtime_checkable
+class CacheableEnvBase(EnvBase, Protocol):
     async def get_cache_assets(self) -> Iterable[str]: ...
 
     @abstractmethod
@@ -277,12 +301,10 @@ async def cache_assets(self) -> None:
 
     @property
     def cache_path(self) -> Path:
-        assert isinstance(self, EnvBase)
         return self._cache_prefix
 
     async def remove_cache(self) -> None:
         """Remove the cached environment assets."""
-        assert isinstance(self, EnvBase)
         for asset in await self.get_cache_assets():
             asset_path = self.cache_path / asset
             if asset_path.exists():
@@ -297,7 +319,8 @@ async def remove_cache(self) -> None:
                     )
 
 
-class DeployableEnvBase(ABC):
+@runtime_checkable
+class DeployableEnvBase(EnvBase, Protocol):
     @abstractmethod
     def is_deployment_path_portable(self) -> bool:
         """Return whether the deployment path matters for the environment, i.e.
@@ -325,7 +348,6 @@ def record_deployment_hash(self, hash_object) -> None:
         deployment is senstivive to the path (e.g. in case of conda, which patches
         the RPATH in binaries).
         """
-        assert isinstance(self, EnvBase)
         self.record_hash(hash_object)
         if not self.is_deployment_path_portable():
             hash_object.update(str(self._deployment_prefix).encode())
@@ -337,24 +359,21 @@ def remove(self) -> None:
 
     def managed_remove(self) -> None:
         """Remove the deployed environment, handling exceptions."""
-        assert isinstance(self, EnvBase)
         try:
             self.remove()
         except Exception as e:
             raise WorkflowError(f"Removal of {self.spec} failed: {e}")
 
     async def managed_deploy(self) -> None:
-        assert isinstance(self, EnvBase)
         try:
             await self.deploy()
         except Exception as e:
             raise WorkflowError(f"Deployment of {self.spec} failed: {e}")
 
     def deployment_hash(self) -> str:
-        assert isinstance(self, EnvBase)
         return self._managed_generic_hash("deployment_hash")
 
     @property
     def deployment_path(self) -> Path:
-        assert isinstance(self, EnvBase) and self._deployment_prefix is not None
+        assert self._deployment_prefix is not None
         return self._deployment_prefix / self.deployment_hash()

snakemake_interface_software_deployment_plugins/tests.py

@@ -90,11 +90,13 @@ def test_deploy(self, tmp_path):
         cmd = env.managed_decorate_shellcmd(self.get_test_cmd())
         assert sp.run(cmd, shell=True, executable=self.shell_executable).returncode == 0
 
-    def test_archive(self, tmp_path):
+    def test_cache(self, tmp_path):
         env = self._get_env(tmp_path)
-        if not isinstance(env, CacheableEnvBase):
+        if not env.is_cacheable():
             pytest.skip("Environment either not deployable or not cacheable.")
 
+        assert isinstance(env, CacheableEnvBase)
+
         asyncio.run(env.cache_assets())
 
         self._deploy(env, tmp_path)
@@ -103,9 +105,11 @@ def test_archive(self, tmp_path):
 
     def test_pin(self, tmp_path):
         env = self._get_env(tmp_path)
-        if not isinstance(env, PinnableEnvBase):
+        if not env.is_pinnable():
             pytest.skip("Environment is not pinnable.")
 
+        assert isinstance(env, PinnableEnvBase)
+
         asyncio.run(env.pin())
         assert env.pinfile.exists()
         print("Pinfile content:", env.pinfile.read_text(), sep="\n")
@@ -164,9 +168,10 @@ def _get_env(self, tmp_path) -> EnvBase:
             pinfile_prefix=pinfile_prefix,
         )
 
-    def _deploy(self, env: DeployableEnvBase, tmp_path):
-        if not isinstance(env, DeployableEnvBase):
+    def _deploy(self, env: EnvBase, tmp_path):
+        if not env.is_deployable():
             pytest.skip("Environment is not deployable.")
 
+        assert isinstance(env, DeployableEnvBase)
         asyncio.run(env.deploy())
         assert any((tmp_path / "deployments").iterdir())