Add prompt capability to dev_build drop-in

lib/spack/spack/build_environment.py

@@ -822,7 +822,7 @@ def load_external_modules(pkg):
             load_module(external_module)
 
 
-def setup_package(pkg, dirty, context: Context = Context.BUILD):
+def setup_package(pkg, dirty, context: Context = Context.BUILD, interactive: bool = False):
     """Execute all environment setup routines."""
     if context not in (Context.BUILD, Context.TEST):
         raise ValueError(f"'context' must be Context.BUILD or Context.TEST - got {context}")
@@ -874,6 +874,8 @@ def setup_package(pkg, dirty, context: Context = Context.BUILD):
 
     load_external_modules(pkg)
 
+    env_mods.set("SPACK_BUILD_ENV", f"{pkg.spec.name}-{pkg.spec.dag_hash()}")
+
     # Make sure nothing's strange about the Spack environment.
     validate(env_mods, tty.warn)
     env_mods.apply_modifications()

lib/spack/spack/cmd/common/env_utility.py

@@ -26,6 +26,7 @@ def setup_parser(subparser):
     subparser.add_argument(
         "--pickle", metavar="FILE", help="dump a pickled source-able environment to FILE"
     )
+    subparser.add_argument("--dive", action="store_true", help="dive into the build-env in a subshell")
     subparser.add_argument(
         "spec",
         nargs=argparse.REMAINDER,
@@ -135,6 +136,9 @@ def emulate_env_utility(cmd_name, context: Context, args):
         tty.msg("Pickling a source-able environment to {0}".format(args.pickle))
         pickle_environment(args.pickle)
 
+    if args.dive:
+        os.execvp(cmd[0], [cmd[0]])
+
     if cmd:
         # Execute the command with the new environment
         os.execvp(cmd[0], cmd)

lib/spack/spack/cmd/dev_build.py

@@ -13,6 +13,7 @@
 import spack.cmd.common.arguments
 import spack.config
 import spack.environment as ev
+import spack.prompt
 import spack.repo
 from spack.cmd.common import arguments
 from spack.installer import PackageInstaller
@@ -65,6 +66,12 @@ def setup_parser(subparser):
         default=None,
         help="drop into a build environment in a new shell, e.g., bash",
     )
+    subparser.add_argument(
+        "-p",
+        "--prompt",
+        action="store_true",
+        help="change the prompt when droping into the build-env",
+    )
     subparser.add_argument(
         "--test",
         default=None,
@@ -162,5 +169,8 @@ def dev_build(self, args):
 
     # drop into the build environment of the package?
     if args.shell is not None:
-        spack.build_environment.setup_package(spec.package, dirty=False)
+        mods = spack.build_environment.setup_package(spec.package, dirty=False)
+        if args.prompt:
+            mods.extend(spack.prompt.prompt_modifications(f"{spec.name}-build-env", args.shell))
+        mods.apply_modifications()
         os.execvp(args.shell, [args.shell])

lib/spack/spack/environment/shell.py

@@ -3,13 +3,12 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 import os
-import textwrap
 from typing import Optional
 
 import llnl.util.tty as tty
-from llnl.util.tty.color import colorize
 
 import spack.environment as ev
+import spack.prompt
 import spack.repo
 import spack.store
 from spack.util.environment import EnvironmentModifications
@@ -24,66 +23,32 @@ def activate_header(env, shell, prompt=None, view: Optional[str] = None):
         if view:
             cmds += "setenv SPACK_ENV_VIEW %s;\n" % view
         cmds += 'alias despacktivate "spack env deactivate";\n'
-        if prompt:
-            cmds += "if (! $?SPACK_OLD_PROMPT ) "
-            cmds += 'setenv SPACK_OLD_PROMPT "${prompt}";\n'
-            cmds += 'set prompt="%s ${prompt}";\n' % prompt
     elif shell == "fish":
-        if "color" in os.getenv("TERM", "") and prompt:
-            prompt = colorize("@G{%s} " % prompt, color=True)
-
         cmds += "set -gx SPACK_ENV %s;\n" % env.path
         if view:
             cmds += "set -gx SPACK_ENV_VIEW %s;\n" % view
         cmds += "function despacktivate;\n"
         cmds += "   spack env deactivate;\n"
         cmds += "end;\n"
-        #
-        # NOTE: We're not changing the fish_prompt function (which is fish's
-        # solution to the PS1 variable) here. This is a bit fiddly, and easy to
-        # screw up => spend time reasearching a solution. Feedback welcome.
-        #
     elif shell == "bat":
         # TODO: Color
         cmds += 'set "SPACK_ENV=%s"\n' % env.path
         if view:
             cmds += 'set "SPACK_ENV_VIEW=%s"\n' % view
         # TODO: despacktivate
-        # TODO: prompt
     elif shell == "pwsh":
         cmds += "$Env:SPACK_ENV='%s'\n" % env.path
         if view:
             cmds += "$Env:SPACK_ENV_VIEW='%s'\n" % view
     else:
-        bash_color_prompt = colorize(f"@G{{{prompt}}}", color=True, enclose=True)
-        zsh_color_prompt = colorize(f"@G{{{prompt}}}", color=True, enclose=False, zsh=True)
-
         cmds += "export SPACK_ENV=%s;\n" % env.path
         if view:
             cmds += "export SPACK_ENV_VIEW=%s;\n" % view
         cmds += "alias despacktivate='spack env deactivate';\n"
+
     if prompt:
-            cmds += textwrap.dedent(
+        cmds += spack.prompt.custom_prompt(prompt, shell)
-                rf"""
+
-                if [ -z ${{SPACK_OLD_PS1+x}} ]; then
-                    if [ -z ${{PS1+x}} ]; then
-                        PS1='$$$$';
-                    fi;
-                    export SPACK_OLD_PS1="${{PS1}}";
-                fi;
-                if [ -n "${{TERM:-}}" ] && [ "${{TERM#*color}}" != "${{TERM}}" ] && \
-                   [ -n "${{BASH:-}}" ];
-                then
-                    export PS1="{bash_color_prompt} ${{PS1}}";
-                elif [ -n "${{TERM:-}}" ] && [ "${{TERM#*color}}" != "${{TERM}}" ] && \
-                     [ -n "${{ZSH_NAME:-}}" ];
-                then
-                    export PS1="{zsh_color_prompt} ${{PS1}}";
-                else
-                    export PS1="{prompt} ${{PS1}}";
-                fi
-                """
-            ).lstrip("\n")
     return cmds
 
 

lib/spack/spack/prompt.py

@@ -0,0 +1,73 @@
+# Copyright 2013-2024 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+import os
+import textwrap
+from llnl.util.tty.color import colorize
+
+from spack.util.environment import EnvironmentModifications
+
+
+def prompt_modifications(prompt, shell, env=os.environ):
+    mods = EnvironmentModifications()
+
+    if shell == "fish" or shell == "pwsh" or shell == "bat":
+        # requires a function and can't be set with os.environ
+        pass
+    elif shell == "csh":
+        mods.set("SPACK_OLD_PROMPT", env.get("prompt", None))
+        mods.set("prompt", prompt)
+    else:
+        mods.set('SPACK_OLD_PS1', env.get('PS1', '$$$$'))
+        if 'TERM' in env and 'color' in env['TERM']:
+            if 'BASH' in env:
+                bash_color_prompt = colorize(f"@G{{{prompt}}}", color=True, enclose=True)
+                mods.set('PS1', f"{bash_color_prompt} {env.get('PS1','$ ')}")
+            else:
+                zsh_color_prompt = colorize(f"@G{{{prompt}}}", color=True, enclose=False, zsh=True)
+                mods.set('PS1', f"{zsh_color_prompt} {env.get('PS1', '$ ')}")
+        else:
+            mods.set('PS1', f"{prompt} {env.get('PS1', '$ ')}")
+
+    return mods
+
+
+def custom_prompt(prompt, shell):
+    cmds = ""
+    if shell == "csh":
+        cmds += "if (! $?SPACK_OLD_PROMPT ) "
+        cmds += 'setenv SPACK_OLD_PROMPT "${prompt}";\n'
+        cmds += 'set prompt="%s ${prompt}";\n' % prompt
+    elif shell == "fish":
+        if "color" in os.getenv("TERM", ""):
+            prompt = colorize(f"@G{prompt} " % prompt, color=True)
+    elif shell == "bat" or shell == "pwsh":
+        # TODO
+        pass
+    else:
+        bash_color_prompt = colorize(f"@G{{{prompt}}}", color=True, enclose=True)
+        zsh_color_prompt = colorize(f"@G{{{prompt}}}", color=True, enclose=False, zsh=True)
+        cmds += textwrap.dedent(
+            rf"""
+            if [ -z ${{SPACK_OLD_PS1+x}} ]; then
+                if [ -z ${{PS1+x}} ]; then
+                    PS1='$$$$';
+                fi;
+                export SPACK_OLD_PS1="${{PS1}}";
+            fi;
+            if [ -n "${{TERM:-}}" ] && [ "${{TERM#*color}}" != "${{TERM}}" ] && \
+               [ -n "${{BASH:-}}" ];
+            then
+                export PS1="{bash_color_prompt} ${{PS1}}";
+            elif [ -n "${{TERM:-}}" ] && [ "${{TERM#*color}}" != "${{TERM}}" ] && \
+                 [ -n "${{ZSH_NAME:-}}" ];
+            then
+                export PS1="{zsh_color_prompt} ${{PS1}}";
+            else
+                export PS1="{prompt} ${{PS1}}";
+            fi
+            """
+        ).lstrip("\n")
+    return cmds
+