Type checking focused around element.py and shell testing

src/buildstream/_artifact.py

@@ -26,17 +26,24 @@
 
 import os
 import tempfile
-from typing import Dict, Tuple
+from typing import Dict, Tuple, Optional, TYPE_CHECKING
 
+from ._cas.cascache import CASCache
+from .sandbox.sandbox import Sandbox
+from .storage.directory import Directory
 from ._protos.buildstream.v2.artifact_pb2 import Artifact as ArtifactProto
 from . import _yaml
 from . import utils
-from .node import Node
+from .node import Node, MappingNode
 from .types import _Scope
 from .storage._casbaseddirectory import CasBasedDirectory
 from .sandbox._config import SandboxConfig
 from ._variables import Variables
 
+if TYPE_CHECKING:
+    from .element import Element
+    from ._context import Context
+
 
 # An Artifact class to abstract artifact operations
 # from the Element class
@@ -52,22 +59,33 @@ class Artifact:
 
     version = 2
 
-    def __init__(self, element, context, *, strong_key=None, strict_key=None, weak_key=None):
-        self._element = element
-        self._context = context
-        self._cache_key = strong_key
-        self._strict_key = strict_key
-        self._weak_cache_key = weak_key
-        self._artifactdir = context.artifactdir
-        self._cas = context.get_cascache()
-        self._tmpdir = context.tmpdir
-        self._proto = None
-
-        self._metadata_keys = None  # Strong, strict and weak key tuple extracted from the artifact
-        self._metadata_dependencies = None  # Dictionary of dependency strong keys from the artifact
-        self._metadata_workspaced = None  # Boolean of whether it's a workspaced artifact
-        self._metadata_workspaced_dependencies = None  # List of which dependencies are workspaced from the artifact
-        self._cached = None  # Boolean of whether the artifact is cached
+    def __init__(
+        self,
+        element: "Element",
+        context: "Context",
+        *,
+        strong_key: Optional[str] = None,
+        strict_key: Optional[str] = None,
+        weak_key: Optional[str] = None,
+    ):
+        self._element: "Element" = element
+        self._context: "Context" = context
+        self._cache_key: Optional[str] = strong_key
+        self._strict_key: Optional[str] = strict_key
+        self._weak_cache_key: Optional[str] = weak_key
+        self._artifactdir: Optional[str] = context.artifactdir
+        self._cas: CASCache = context.get_cascache()
+        self._tmpdir: Optional[str] = context.tmpdir
+        self._proto: Optional[ArtifactProto] = None
+
+        self._metadata_keys: Optional[tuple[str, str, str]] = (
+            None  # Strong, strict and weak key tuple extracted from the artifact
+        )
+        self._metadata_workspaced: Optional[bool] = None  # Boolean of whether it's a workspaced artifact
+        self._metadata_workspaced_dependencies: Optional[list[str]] = (
+            None  # List of which dependencies are workspaced from the artifact
+        )
+        self._cached: Optional[bool] = None  # Boolean of whether the artifact is cached
 
     # strong_key():
     #
@@ -76,7 +94,8 @@ def __init__(self, element, context, *, strong_key=None, strict_key=None, weak_k
     # or whether it was the strong key loaded from artifact metadata.
     #
     @property
-    def strong_key(self) -> str:
+    def strong_key(self) -> Optional[str]:
+        key: str | None
         if self.cached():
             key, _, _ = self.get_metadata_keys()
         else:
@@ -91,7 +110,8 @@ def strong_key(self) -> str:
     # or whether it was the strict key loaded from artifact metadata.
     #
     @property
-    def strict_key(self) -> str:
+    def strict_key(self) -> Optional[str]:
+        key: str | None
         if self.cached():
             _, key, _ = self.get_metadata_keys()
         else:
@@ -106,7 +126,8 @@ def strict_key(self) -> str:
     # or whether it was the weak key loaded from artifact metadata.
     #
     @property
-    def weak_key(self) -> str:
+    def weak_key(self) -> Optional[str]:
+        key: str | None
         if self.cached():
             _, _, key = self.get_metadata_keys()
         else:
@@ -121,7 +142,7 @@ def weak_key(self) -> str:
     # Returns:
     #    (Directory): The virtual directory object
     #
-    def get_files(self):
+    def get_files(self) -> Directory:
         files_digest = self._get_field_digest("files")
         return CasBasedDirectory(self._cas, digest=files_digest)
 
@@ -235,6 +256,9 @@ def cache(
         artifact.build_error_details = "" if not buildresult[2] else buildresult[2]
 
         # Store keys
+        assert self._cache_key, "Key should be ready by now"
+        assert self._strict_key, "Key should be ready by now"
+        assert self._weak_cache_key, "Key should be ready by now"
         artifact.strong_key = self._cache_key
         artifact.strict_key = self._strict_key
         artifact.weak_key = self._weak_cache_key
@@ -292,7 +316,7 @@ def cache(
             digests = self._cas.add_objects(paths=[entry[0] for entry in files_to_capture])
             # add_objects() should guarantee this.
             # `zip(..., strict=True)` could be used in Python 3.10+
-            assert len(files_to_capture) == len(digests)
+            assert len(files_to_capture) == len(digests), "files_to_capture and digests should be the same length"
             for entry, digest in zip(files_to_capture, digests):
                 entry[1].CopyFrom(digest)
 
@@ -336,6 +360,7 @@ def cache(
                 digest = artifact.buildsandbox.subsandbox_digests.add()
                 digest.CopyFrom(vdir._get_digest())
 
+        assert self._artifactdir, "An artifact dir is required at this point"
         os.makedirs(os.path.dirname(os.path.join(self._artifactdir, element.get_artifact_name())), exist_ok=True)
         keys = utils._deduplicate([self._cache_key, self._weak_cache_key])
         for key in keys:
@@ -370,7 +395,7 @@ def cached_buildroot(self):
     # Returns:
     #     (bool): True if artifact was created with buildroot
     #
-    def buildroot_exists(self):
+    def buildroot_exists(self) -> bool:
 
         artifact = self._get_proto()
         return bool(str(artifact.buildroot))
@@ -386,7 +411,7 @@ def buildroot_exists(self):
     #             missing expected buildtree. Note this only confirms
     #             if a buildtree is present, not its contents.
     #
-    def cached_buildtree(self):
+    def cached_buildtree(self) -> bool:
 
         buildtree_digest = self._get_field_digest("buildtree")
         if buildtree_digest:
@@ -430,7 +455,7 @@ def cached_sources(self):
     # Returns:
     #    (dict): The artifacts cached public data
     #
-    def load_public_data(self):
+    def load_public_data(self) -> MappingNode:
 
         # Load the public data from the artifact
         artifact = self._get_proto()
@@ -546,7 +571,7 @@ def get_metadata_keys(self) -> Tuple[str, str, str]:
     # Returns:
     #    (bool): Whether the given artifact was workspaced
     #
-    def get_metadata_workspaced(self):
+    def get_metadata_workspaced(self) -> bool:
 
         if self._metadata_workspaced is not None:
             return self._metadata_workspaced
@@ -565,7 +590,7 @@ def get_metadata_workspaced(self):
     # Returns:
     #    (list): List of which dependencies are workspaced
     #
-    def get_metadata_workspaced_dependencies(self):
+    def get_metadata_workspaced_dependencies(self) -> list[str]:
 
         if self._metadata_workspaced_dependencies is not None:
             return self._metadata_workspaced_dependencies
@@ -645,7 +670,7 @@ def query_cache(self):
     #     (bool): Whether artifact is in local cache
     #
     def cached(self, *, buildtree=False):
-        assert self._cached is not None
+        assert self._cached is not None, "_cached should have been initialised before calling this method"
         ret = self._cached
         if buildtree:
             ret = ret and (self.cached_buildtree() or not self.buildtree_exists())
@@ -671,7 +696,7 @@ def cached_logs(self):
     #
     def set_cached(self):
         self._proto = self._load_proto()
-        assert self._proto
+        assert self._proto, "We expect a value is returned by _load_proto()"
         self._cached = True
 
     # pull()
@@ -700,7 +725,7 @@ def pull(self, *, pull_buildtrees):
 
         return True
 
-    def configure_sandbox(self, sandbox):
+    def configure_sandbox(self, sandbox: Sandbox):
         artifact = self._get_proto()
 
         if artifact.HasField("buildsandbox") and artifact.buildsandbox.environment:
@@ -722,13 +747,13 @@ def configure_sandbox(self, sandbox):
     #  load_proto()
     #
     # Returns:
-    #     (Artifact): Artifact proto
+    #     (ArtifactProto): Artifact proto
     #
-    def _load_proto(self):
+    def _load_proto(self) -> Optional["ArtifactProto"]:
         key = self.get_extract_key()
-
+        assert self._artifactdir, "Must have and artifact dir to load proto"
         proto_path = os.path.join(self._artifactdir, self._element.get_artifact_name(key=key))
-        artifact = ArtifactProto()
+        artifact: ArtifactProto = ArtifactProto()
         try:
             with open(proto_path, mode="r+b") as f:
                 artifact.ParseFromString(f.read())

src/buildstream/_artifactcache.py

@@ -14,15 +14,20 @@
 #  Authors:
 #        Tristan Maat <tristan.maat@codethink.co.uk>
 
+
 import os
+from typing import TYPE_CHECKING
 
 from ._assetcache import AssetCache
 from ._cas.casremote import BlobNotFound
-from ._exceptions import ArtifactError, AssetCacheError, CASError, CASRemoteError
+from ._exceptions import ArtifactError, AssetCacheError, CASError, CASRemoteError, BstError
 from ._protos.buildstream.v2 import artifact_pb2
-
 from . import utils
 
+if TYPE_CHECKING:
+    from . import Element
+    from ._context import Context
+
 REMOTE_ASSET_ARTIFACT_URN_TEMPLATE = "urn:fdc:buildstream.build:2020:artifact:{}"
 
 
@@ -32,11 +37,12 @@
 #     context (Context): The BuildStream context
 #
 class ArtifactCache(AssetCache):
-    def __init__(self, context):
+    def __init__(self, context: "Context"):
         super().__init__(context)
 
         # create artifact directory
         self._basedir = context.artifactdir
+        assert self._basedir is not None, "Must have a base dir"
         os.makedirs(self._basedir, exist_ok=True)
 
     # preflight():
@@ -57,9 +63,9 @@ def preflight(self):
     #
     # Returns: True if the artifact is in the cache, False otherwise
     #
-    def contains(self, element, key):
+    def contains(self, element: "Element", key: str):
         ref = element.get_artifact_name(key)
-
+        assert self._basedir is not None, "Must have a base directory"
         return os.path.exists(os.path.join(self._basedir, ref))
 
     # list_artifacts():
@@ -151,7 +157,7 @@ def push(self, element, artifact):
     # Returns:
     #   (bool): True if pull was successful, False if artifact was not available
     #
-    def pull(self, element, key, *, pull_buildtrees=False):
+    def pull(self, element: "Element", key: str, *, pull_buildtrees: bool = False) -> bool:
         artifact_digest = None
         display_key = key[: self.context.log_key_length]
         project = element._get_project()
@@ -161,7 +167,7 @@ def pull(self, element, key, *, pull_buildtrees=False):
 
         index_remotes, storage_remotes = self.get_remotes(project.name, False)
 
-        errors = []
+        errors: list[BstError] = []
         # Start by pulling our artifact proto, so that we know which
         # blobs to pull
         for remote in index_remotes:
@@ -192,22 +198,24 @@ def pull(self, element, key, *, pull_buildtrees=False):
 
         errors = []
         # If we do, we can pull it!
-        for remote in storage_remotes:
-            remote.init()
+        for storage_remote in storage_remotes:
+            storage_remote.init()
             try:
-                element.status("Pulling data for artifact {} <- {}".format(display_key, remote))
+                element.status("Pulling data for artifact {} <- {}".format(display_key, storage_remote))
 
-                if self._pull_artifact_storage(element, key, artifact_digest, remote, pull_buildtrees=pull_buildtrees):
-                    element.info("Pulled artifact {} <- {}".format(display_key, remote))
+                if self._pull_artifact_storage(
+                    element, key, artifact_digest, storage_remote, pull_buildtrees=pull_buildtrees
+                ):
+                    element.info("Pulled artifact {} <- {}".format(display_key, storage_remote))
                     return True
 
-                element.info("Remote ({}) does not have artifact {} cached".format(remote, display_key))
+                element.info("Remote ({}) does not have artifact {} cached".format(storage_remote, display_key))
             except BlobNotFound as e:
-                # Not all blobs are available on this remote
-                element.info("Remote cas ({}) does not have blob {} cached".format(remote, e.blob))
+                # Not all blobs are available on this storage_remote
+                element.info("Remote cas ({}) does not have blob {} cached".format(storage_remote, e.blob))
                 continue
             except CASError as e:
-                element.warn("Could not pull from remote {}: {}".format(remote, e))
+                element.warn("Could not pull from remote {}: {}".format(storage_remote, e))
                 errors.append(e)
 
         if errors:
@@ -235,7 +243,7 @@ def link_key(self, element, oldkey, newkey):
         if oldref == newref:
             # The two refs are identical, nothing to do
             return
-
+        assert self._basedir, "base dir should be set"
         utils.safe_link(os.path.join(self._basedir, oldref), os.path.join(self._basedir, newref))
 
     # check_remotes_for_element()
@@ -248,7 +256,7 @@ def link_key(self, element, oldkey, newkey):
     # Returns:
     #    (bool): True if the element is available remotely
     #
-    def check_remotes_for_element(self, element):
+    def check_remotes_for_element(self, element: "Element") -> bool:
         project = element._get_project()
         index_remotes, _ = self.get_remotes(project.name, False)
 
@@ -421,7 +429,8 @@ def _pull_artifact_storage(self, element, key, artifact_digest, remote, pull_bui
             with self.cas.open(artifact_digest, "rb") as f:
                 artifact.ParseFromString(f.read())
 
-            # Write the artifact proto to cache
+            # Write the artifact proto to cache#
+            assert self._basedir, "Must have a base directory"
             artifact_path = os.path.join(self._basedir, artifact_name)
             os.makedirs(os.path.dirname(artifact_path), exist_ok=True)
             with utils.save_file_atomic(artifact_path, mode="wb") as f: