Move shell element loading logic from Cli to Stream

3 files changed, 73 insertions, 69 deletions

diff --git a/src/buildstream/_frontend/cli.py b/src/buildstream/_frontend/cli.py
index 3010583ab..3f804c2af 100644
--- a/src/buildstream/_frontend/cli.py
+++ b/src/buildstream/_frontend/cli.py
@@ -620,78 +620,19 @@ def shell(app, element, mount, isolate, build_, cli_buildtree, pull_, command):
     # Buildtree can only be used with build shells
     if cli_buildtree != "never":
         build_ = True
+    else:
+        cli_buildtree = None
 
     scope = _Scope.BUILD if build_ else _Scope.RUN
 
-    # We may need to fetch dependency artifacts if we're pulling the artifact
-    selection = _PipelineSelection.ALL if pull_ else _PipelineSelection.NONE
-    use_buildtree = None
-
     with app.initialized():
         if not element:
             element = app.project.get_default_target()
             if not element:
                 raise AppError('Missing argument "ELEMENT".')
 
-        elements = app.stream.load_selection((element,), selection=selection, use_artifact_config=True)
-
-        # last one will be the element we want to stage, previous ones are
-        # elements to try and pull
-        element = elements[-1]
-        pull_dependencies = elements[:-1] if pull_ else None
-
         mounts = [HostMount(path, host_path) for host_path, path in mount]
 
-        artifact_is_cached = element._cached()
-        buildtree_is_cached = element._cached_buildtree()
-        buildtree_exists = element._buildtree_exists()
-        can_attempt_pull = app.context.pull_buildtrees and pull_
-
-        if cli_buildtree in ("always", "try"):
-            if buildtree_is_cached:
-                use_buildtree = cli_buildtree
-            # If element is already cached, we can check the proto to see if the buildtree existed
-            elif artifact_is_cached:
-                if not buildtree_exists:
-                    if cli_buildtree == "always":
-                        # Exit early if it won't be possible to even fetch a buildtree with always option
-                        raise AppError("Artifact was created without buildtree, unable to launch shell with it")
-                    click.echo(
-                        "WARNING: Artifact created without buildtree, shell will be loaded without it", err=True
-                    )
-                elif can_attempt_pull:
-                    use_buildtree = cli_buildtree
-                    click.echo(
-                        "WARNING: buildtree is not cached locally but did exist, will attempt to pull from available remotes",
-                        err=True,
-                    )
-                else:
-                    if cli_buildtree == "always":
-                        # Exit early if it won't be possible to perform a fetch as pull semantics aren't present
-                        raise AppError(
-                            "Artifact has a buildtree but it isn't cached. Can be retried with --pull and pull-buildtrees configured"
-                        )
-                    click.echo("WARNING: buildtree is not cached locally, shell will be loaded without it", err=True)
-            # If element isn't cached at all, we can't check the proto to see if it existed so can't exit early
-            elif can_attempt_pull:
-                use_buildtree = cli_buildtree
-                if use_buildtree == "always":
-                    click.echo(
-                        "WARNING: Element is not cached so buildtree status unknown, will attempt to pull from available remotes",
-                        err=True,
-                    )
-            else:
-                if cli_buildtree == "always":
-                    # Exit early as there is no buildtree locally & can_attempt_pull is False
-                    raise AppError(
-                        "Artifact not cached locally. Can be retried with --pull and pull-buildtrees configured"
-                    )
-                click.echo("WARNING: buildtree is not cached locally, shell will be loaded without it", err=True)
-
-        # Raise warning if the element is cached in a failed state
-        if use_buildtree and element._cached_failure():
-            click.echo("WARNING: using a buildtree from a failed build.", err=True)
-
         try:
             exitcode = app.stream.shell(
                 element,
@@ -700,8 +641,8 @@ def shell(app, element, mount, isolate, build_, cli_buildtree, pull_, command):
                 mounts=mounts,
                 isolate=isolate,
                 command=command,
-                usebuildtree=use_buildtree,
-                pull_dependencies=pull_dependencies,
+                usebuildtree=cli_buildtree,
+                pull_=pull_,
             )
         except BstError as e:
             raise AppError("Error launching shell: {}".format(e), detail=e.detail) from e
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 19e371337..766c473ae 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -167,14 +167,14 @@ class Stream:
     # Run a shell
     #
     # Args:
-    #    element (Element): An Element object to run the shell for
+    #    element (str): The name of the element to run the shell for
     #    scope (_Scope): The scope for the shell (_Scope.BUILD or _Scope.RUN)
     #    prompt (function): A function to return the prompt to display in the shell
     #    mounts (list of HostMount): Additional directories to mount into the sandbox
     #    isolate (bool): Whether to isolate the environment like we do in builds
     #    command (list): An argv to launch in the sandbox, or None
     #    usebuildtree (str): Whether to use a buildtree as the source, given cli option
-    #    pull_dependencies ([Element]|None): Elements to attempt to pull
+    #    pull_ (bool): Whether to attempt to pull missing or incomplete artifacts
     #    unique_id: (str): Whether to use a unique_id to load an Element instance
     #
     # Returns:
@@ -190,13 +190,76 @@ class Stream:
         isolate=False,
         command=None,
         usebuildtree=None,
-        pull_dependencies=None,
+        pull_=False,
         unique_id=None
     ):
 
         # Load the Element via the unique_id if given
         if unique_id and element is None:
             element = Plugin._lookup(unique_id)
+        else:
+            # We may need to fetch dependency artifacts if we're pulling the artifact
+            selection = _PipelineSelection.ALL if pull_ else _PipelineSelection.NONE
+
+            elements = self.load_selection((element,), selection=selection, use_artifact_config=True)
+
+            # last one will be the element we want to stage, previous ones are
+            # elements to try and pull
+            element = elements[-1]
+            pull_dependencies = elements[:-1] if pull_ else None
+
+            artifact_is_cached = element._cached()
+            buildtree_is_cached = element._cached_buildtree()
+            buildtree_exists = element._buildtree_exists()
+            can_attempt_pull = self._context.pull_buildtrees and pull_
+
+            if usebuildtree:
+                if buildtree_is_cached:
+                    pass
+                # If element is already cached, we can check the proto to see if the buildtree existed
+                elif artifact_is_cached:
+                    if not buildtree_exists:
+                        if usebuildtree == "always":
+                            # Exit early if it won't be possible to even fetch a buildtree with always option
+                            raise StreamError("Artifact was created without buildtree, unable to launch shell with it")
+                        usebuildtree = None
+                        self._message(
+                            MessageType.WARN, "Artifact created without buildtree, shell will be loaded without it"
+                        )
+                    elif can_attempt_pull:
+                        self._message(
+                            MessageType.WARN,
+                            "buildtree is not cached locally but did exist, will attempt to pull from available remotes",
+                        )
+                    else:
+                        if usebuildtree == "always":
+                            # Exit early if it won't be possible to perform a fetch as pull semantics aren't present
+                            raise StreamError(
+                                "Artifact has a buildtree but it isn't cached. Can be retried with --pull and pull-buildtrees configured"
+                            )
+                        usebuildtree = None
+                        self._message(
+                            MessageType.WARN, "buildtree is not cached locally, shell will be loaded without it"
+                        )
+                # If element isn't cached at all, we can't check the proto to see if it existed so can't exit early
+                elif can_attempt_pull:
+                    if usebuildtree == "always":
+                        self._message(
+                            MessageType.WARN,
+                            "Element is not cached so buildtree status unknown, will attempt to pull from available remotes",
+                        )
+                else:
+                    if usebuildtree == "always":
+                        # Exit early as there is no buildtree locally & can_attempt_pull is False
+                        raise StreamError(
+                            "Artifact not cached locally. Can be retried with --pull and pull-buildtrees configured"
+                        )
+                    usebuildtree = None
+                    self._message(MessageType.WARN, "buildtree is not cached locally, shell will be loaded without it")
+
+            # Raise warning if the element is cached in a failed state
+            if usebuildtree and element._cached_failure():
+                self._message(MessageType.WARN, "using a buildtree from a failed build.")
 
         missing_deps = [dep for dep in self._pipeline.dependencies([element], scope) if not dep._cached()]
         if missing_deps:
diff --git a/tests/integration/shellbuildtrees.py b/tests/integration/shellbuildtrees.py
index 0d80c1640..bb3a758e1 100644
--- a/tests/integration/shellbuildtrees.py
+++ b/tests/integration/shellbuildtrees.py
@@ -131,7 +131,7 @@ def test_buildtree_from_failure(cli_integration, datafiles):
         project=project, args=["shell", "--build", element_name, "--use-buildtree", "always", "--", "cat", "test"]
     )
     res.assert_success()
-    assert "WARNING: using a buildtree from a failed build" in res.stderr
+    assert "WARNING using a buildtree from a failed build" in res.stderr
     assert "Hi" in res.output
 
 
@@ -175,7 +175,7 @@ def test_buildtree_from_failure_option_always(cli_integration, tmpdir, datafiles
         project=project, args=["shell", "--build", element_name, "--use-buildtree", "always", "--", "cat", "test"]
     )
     res.assert_success()
-    assert "WARNING: using a buildtree from a failed build" in res.stderr
+    assert "WARNING using a buildtree from a failed build" in res.stderr
     assert "Hi" in res.output
 
 
@@ -257,7 +257,7 @@ def test_buildtree_options(cli, tmpdir, datafiles):
         )
         assert "Hi" not in res.output
         assert "Attempting to fetch missing artifact buildtrees" not in res.stderr
-        assert "WARNING: buildtree is not cached locally, shell will be loaded without it" in res.stderr
+        assert "WARNING buildtree is not cached locally, shell will be loaded without it" in res.stderr
 
         # Check correctly handling the lack of buildtree, with 'try' attempting and succeeding
         # to pull the buildtree as the user context allow the pulling of buildtrees and it is