_stream.py: Remove separate handling of ArtifactElements

ArtifactElement inherits Element, both have an Artifact
object as a member, thus we should not need to handle these
separately.

This change has resulted in introducing a configure_sandbox()
method in ArtifactElement. The method is similar to
BuildElement.configure_sandbox() but does not configure the
sandbox to actually be used for building.

3 files changed, 29 insertions, 22 deletions

diff --git a/src/buildstream/_artifactelement.py b/src/buildstream/_artifactelement.py
index e3bffaf72..48c3d1769 100644
--- a/src/buildstream/_artifactelement.py
+++ b/src/buildstream/_artifactelement.py
@@ -42,6 +42,10 @@ class ArtifactElement(Element):
     # A hash of ArtifactElement by ref
     __instantiated_artifacts = {}   # type: Dict[str, ArtifactElement]
 
+    # ArtifactElement's require this as the sandbox will use a normal
+    # directory when we checkout
+    BST_VIRTUAL_DIRECTORY = True
+
     def __init__(self, context, ref):
         _, element, key = verify_artifact_ref(ref)
 
@@ -126,6 +130,22 @@ class ArtifactElement(Element):
         artifact = self._get_artifact()
         return artifact.get_dependency_refs(deps=scope)
 
+    # configure_sandbox()
+    #
+    # Configure a sandbox for installing artifacts into
+    #
+    # Args:
+    #    sandbox (Sandbox)
+    #
+    def configure_sandbox(self, sandbox):
+        install_root = self.get_variable('install-root')
+
+        # Tell the sandbox to mount the build root and install root
+        sandbox.mark_directory(install_root)
+
+        # Tell sandbox which directory is preserved in the finished artifact
+        sandbox.set_output_directory(install_root)
+
     # Override Element._calculate_cache_key
     def _calculate_cache_key(self, dependencies=None):
         return self._key
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 083dc1c43..a28ca80f6 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -29,7 +29,6 @@ import tempfile
 from contextlib import contextmanager, suppress
 from fnmatch import fnmatch
 
-from ._artifact import Artifact
 from ._artifactelement import verify_artifact_ref, ArtifactElement
 from ._exceptions import StreamError, ImplError, BstError, ArtifactElementError, ArtifactError
 from ._message import Message, MessageType
@@ -551,27 +550,15 @@ class Stream():
             self._enqueue_plan(uncached_elts)
             self._run()
 
-        # Stage deps into a temporary sandbox first
-        if isinstance(target, ArtifactElement):
-            try:
-                key = target._get_cache_key()
-                artifact = Artifact(target, self._context, strong_key=key)
-                virdir = artifact.get_files()
+        try:
+            with target._prepare_sandbox(scope=scope, directory=None,
+                                         integrate=integrate) as sandbox:
+                # Copy or move the sandbox to the target directory
+                virdir = sandbox.get_virtual_directory()
                 self._export_artifact(tar, location, compression, target, hardlinks, virdir)
-            except AttributeError as e:
-                raise ArtifactError("Artifact reference '{}' seems to be invalid. "
-                                    "Note that an Element name can also be used."
-                                    .format(artifact._element.get_artifact_name())) from e
-        else:
-            try:
-                with target._prepare_sandbox(scope=scope, directory=None,
-                                             integrate=integrate) as sandbox:
-                    # Copy or move the sandbox to the target directory
-                    virdir = sandbox.get_virtual_directory()
-                    self._export_artifact(tar, location, compression, target, hardlinks, virdir)
-            except BstError as e:
-                raise StreamError("Error while staging dependencies into a sandbox"
-                                  ": '{}'".format(e), detail=e.detail, reason=e.reason) from e
+        except BstError as e:
+            raise StreamError("Error while staging dependencies into a sandbox"
+                              ": '{}'".format(e), detail=e.detail, reason=e.reason) from e
 
     def _export_artifact(self, tar, location, compression, target, hardlinks, virdir):
         if not tar:
diff --git a/tests/frontend/buildcheckout.py b/tests/frontend/buildcheckout.py
index 6281217b7..72974a620 100644
--- a/tests/frontend/buildcheckout.py
+++ b/tests/frontend/buildcheckout.py
@@ -358,7 +358,7 @@ def test_build_checkout_invalid_ref(datafiles, cli):
     checkout_args = ['artifact', 'checkout', '--deps', 'none', '--tar', checkout, non_existent_artifact]
     result = cli.run(project=project, args=checkout_args)
 
-    assert "Artifact reference '{}' seems to be invalid".format(non_existent_artifact) in result.stderr
+    assert "Error while staging dependencies into a sandbox: 'No artifacts to stage'" in result.stderr
 
 
 @pytest.mark.datafiles(DATA_DIR)