_stream.py: Use callback for shell prompt

3 files changed, 15 insertions, 14 deletions

diff --git a/src/buildstream/_frontend/app.py b/src/buildstream/_frontend/app.py
index b25a421c3..59bc3513b 100644
--- a/src/buildstream/_frontend/app.py
+++ b/src/buildstream/_frontend/app.py
@@ -446,15 +446,16 @@ class App:
     # if they are available in the execution context.
     #
     # Args:
-    #    element_name (str): The element's full name
-    #    element_key (tuple): The element's display key
+    #    element (Element): The element
     #
     # Returns:
     #    (str): The formatted prompt to display in the shell
     #
-    def shell_prompt(self, element_name, element_key):
+    def shell_prompt(self, element):
 
-        _, key, dim = element_key
+        element_name = element._get_full_name()
+
+        _, key, dim = element._get_display_key()
 
         if self.colors:
             prompt = (
@@ -703,10 +704,14 @@ class App:
                 if choice == "shell":
                     click.echo("\nDropping into an interactive shell in the failed build sandbox\n", err=True)
                     try:
-                        unique_id, element_key = element
-                        prompt = self.shell_prompt(full_name, element_key)
+                        unique_id, _ = element
                         self.stream.shell(
-                            None, _Scope.BUILD, prompt, isolate=True, usebuildtree="always", unique_id=unique_id
+                            None,
+                            _Scope.BUILD,
+                            self.shell_prompt,
+                            isolate=True,
+                            usebuildtree="always",
+                            unique_id=unique_id,
                         )
                     except BstError as e:
                         click.echo("Error while attempting to create interactive shell: {}".format(e), err=True)
diff --git a/src/buildstream/_frontend/cli.py b/src/buildstream/_frontend/cli.py
index 24ac6912e..3010583ab 100644
--- a/src/buildstream/_frontend/cli.py
+++ b/src/buildstream/_frontend/cli.py
@@ -640,10 +640,6 @@ def shell(app, element, mount, isolate, build_, cli_buildtree, pull_, command):
         element = elements[-1]
         pull_dependencies = elements[:-1] if pull_ else None
 
-        element_name = element._get_full_name()
-        element_key = element._get_display_key()
-
-        prompt = app.shell_prompt(element_name, element_key)
         mounts = [HostMount(path, host_path) for host_path, path in mount]
 
         artifact_is_cached = element._cached()
@@ -700,7 +696,7 @@ def shell(app, element, mount, isolate, build_, cli_buildtree, pull_, command):
             exitcode = app.stream.shell(
                 element,
                 scope,
-                prompt,
+                app.shell_prompt,
                 mounts=mounts,
                 isolate=isolate,
                 command=command,
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 3b0a308e7..19e371337 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -169,7 +169,7 @@ class Stream:
     # Args:
     #    element (Element): An Element object to run the shell for
     #    scope (_Scope): The scope for the shell (_Scope.BUILD or _Scope.RUN)
-    #    prompt (str): The prompt to display in the shell
+    #    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
@@ -243,7 +243,7 @@ class Stream:
             self._pipeline.assert_sources_cached([element])
 
         return element._shell(
-            scope, mounts=mounts, isolate=isolate, prompt=prompt, command=command, usebuildtree=buildtree
+            scope, mounts=mounts, isolate=isolate, prompt=prompt(element), command=command, usebuildtree=buildtree
         )
 
     # build()