Load artifact refs the same way we load element names

5 files changed, 102 insertions, 6 deletions

diff --git a/src/buildstream/_artifactelement.py b/src/buildstream/_artifactelement.py
index 35ef89574..cfd3c29c8 100644
--- a/src/buildstream/_artifactelement.py
+++ b/src/buildstream/_artifactelement.py
@@ -32,6 +32,9 @@ from .types import Scope
 #    ref (str): The artifact ref
 #
 class ArtifactElement(Element):
+
+    __instantiated_artifacts = {}  # A hash of ArtifactElement by ref
+
     def __init__(self, context, ref):
         _, element, key = verify_artifact_ref(ref)
 
@@ -44,6 +47,52 @@ class ArtifactElement(Element):
 
         super().__init__(context, project, meta, plugin_conf)
 
+    # _new_from_artifact_ref():
+    #
+    # Recursively instantiate a new ArtifactElement instance, and its
+    # dependencies from an artifact ref
+    #
+    # Args:
+    #    ref (String): The artifact ref
+    #    context (Context): The Context object
+    #    task (Task): A task object to report progress to
+    #
+    # Returns:
+    #    (ArtifactElement): A newly created Element instance
+    #
+    @classmethod
+    def _new_from_artifact_ref(cls, ref, context, task=None):
+
+        if ref in cls.__instantiated_artifacts:
+            return cls.__instantiated_artifacts[ref]
+
+        artifact_element = ArtifactElement(context, ref)
+        # XXX: We need to call update state as it is responsible for
+        # initialising an Element/ArtifactElement's Artifact (__artifact)
+        artifact_element._update_state()
+        cls.__instantiated_artifacts[ref] = artifact_element
+
+        for dep_ref in artifact_element.get_dependency_refs(Scope.BUILD):
+            dependency = ArtifactElement._new_from_artifact_ref(dep_ref, context, task)
+            artifact_element._add_build_dependency(dependency)
+
+        return artifact_element
+
+    # _clear_artifact_refs_cache()
+    #
+    # Clear the internal artifact refs cache
+    #
+    # When loading ArtifactElements from artifact refs, we cache already
+    # instantiated ArtifactElements in order to not have to load the same
+    # ArtifactElements twice. This clears the cache.
+    #
+    # It should be called whenever we are done loading all artifacts in order
+    # to save memory.
+    #
+    @classmethod
+    def _clear_artifact_refs_cache(cls):
+        cls.__instantiated_artifacts = {}
+
     # Override Element.get_artifact_name()
     def get_artifact_name(self, key=None):
         return self._ref
diff --git a/src/buildstream/_pipeline.py b/src/buildstream/_pipeline.py
index b6896b3de..7cf4abbe3 100644
--- a/src/buildstream/_pipeline.py
+++ b/src/buildstream/_pipeline.py
@@ -115,6 +115,21 @@ class Pipeline():
 
             return tuple(element_groups)
 
+    # load_artifacts()
+    #
+    # Loads ArtifactElements from target artifacts.
+    #
+    # Args:
+    #    target (list [str]): Target artifacts to load
+    #
+    # Returns:
+    #    (list [ArtifactElement]): A list of ArtifactElement objects
+    #
+    def load_artifacts(self, targets):
+        # XXX: This is not included as part of the "load-pipeline" profiler, we could move
+        #      the profiler to Stream?
+        return self._project.load_artifacts(targets)
+
     # resolve_elements()
     #
     # Resolve element state and cache keys.
diff --git a/src/buildstream/_project.py b/src/buildstream/_project.py
index d3bbc3939..00beebfab 100644
--- a/src/buildstream/_project.py
+++ b/src/buildstream/_project.py
@@ -463,6 +463,38 @@ class Project():
 
         return elements
 
+    # load_artifacts()
+    #
+    # Loads artifacts from target artifact refs
+    #
+    # Args:
+    #    targets (list): Target artifact refs
+    #
+    # Returns:
+    #    (list): A list of loaded ArtifactElement
+    #
+    def load_artifacts(self, targets):
+        with self._context.messenger.simple_task("Loading artifacts") as task:
+            # XXX: Here, we are explicitly checking for refs in the artifactdir
+            #      for two reasons:
+            #          1. The Project, or the Context, do not currently have
+            #             access to the ArtifactCache
+            #          2. The ArtifactCache.contains() method expects an Element
+            #             and a key, not a ref.
+            #
+            artifactdir = self._context.artifactdir
+            artifacts = []
+            for ref in targets:
+                if not os.path.exists(os.path.join(artifactdir, ref)):
+                    raise LoadError("{}\nis not present in the artifact cache ({})".format(ref, artifactdir),
+                                    LoadErrorReason.MISSING_FILE)
+
+                artifacts.append(ArtifactElement._new_from_artifact_ref(ref, self._context, task))
+
+        ArtifactElement._clear_artifact_refs_cache()
+
+        return artifacts
+
     # ensure_fully_loaded()
     #
     # Ensure project has finished loading. At first initialization, a
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 167127cf2..89febc547 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -1172,8 +1172,8 @@ class Stream():
             self._pipeline.load([target_elements, except_targets, track_targets, track_except_targets],
                                 rewritable=rewritable)
 
-        # Obtain the ArtifactElement objects
-        artifacts = [self._project.create_artifact_element(ref) for ref in target_artifacts]
+        # Load all target artifacts
+        artifacts = self._pipeline.load_artifacts(target_artifacts) if target_artifacts else []
 
         # Optionally filter out junction elements
         if ignore_junction_targets:
diff --git a/tests/frontend/buildcheckout.py b/tests/frontend/buildcheckout.py
index b83fc2f3d..98b179b9e 100644
--- a/tests/frontend/buildcheckout.py
+++ b/tests/frontend/buildcheckout.py
@@ -354,11 +354,11 @@ def test_build_checkout_invalid_ref(datafiles, cli):
     assert os.path.isdir(builddir)
     assert not os.listdir(builddir)
 
-    checkout_args = ['artifact', 'checkout', '--deps', 'none', '--tar', checkout,
-                     'test/checkout-deps/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa']
-
+    non_existent_artifact = 'test/checkout-deps/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'
+    checkout_args = ['artifact', 'checkout', '--deps', 'none', '--tar', checkout, non_existent_artifact]
     result = cli.run(project=project, args=checkout_args)
-    assert "seems to be invalid. Note that an Element name can also be used." in result.stderr
+
+    assert "{}\nis not present in the artifact cache".format(non_existent_artifact) in result.stderr
 
 
 @pytest.mark.datafiles(DATA_DIR)