From 2acf90f7b70fcac917d88c3a620505f51df3ef7a Mon Sep 17 00:00:00 2001
From: psakiev <psakiev@sandia.gov>
Date: Mon, 11 Nov 2024 17:05:05 -0700
Subject: [PATCH] WIP unify

---
 lib/spack/spack/build_environment.py      |  2 +-
 lib/spack/spack/cmd/build_env.py          |  4 +++-
 lib/spack/spack/cmd/common/env_utility.py | 14 +++++++++-----
 lib/spack/spack/prompt.py                 | 14 +++++++-------
 4 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/lib/spack/spack/build_environment.py b/lib/spack/spack/build_environment.py
index 14937f2fdd9..3319acca301 100644
--- a/lib/spack/spack/build_environment.py
+++ b/lib/spack/spack/build_environment.py
@@ -805,7 +805,7 @@ def load_external_modules(pkg):
             load_module(external_module)
 
 
-def setup_package(pkg, dirty, context: Context = Context.BUILD, interactive: bool = False):
+def setup_package(pkg, dirty, context: Context = Context.BUILD):
     """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}")
diff --git a/lib/spack/spack/cmd/build_env.py b/lib/spack/spack/cmd/build_env.py
index 2083b4a5b26..89fcd367831 100644
--- a/lib/spack/spack/cmd/build_env.py
+++ b/lib/spack/spack/cmd/build_env.py
@@ -6,7 +6,7 @@
 from spack.context import Context
 
 description = (
-    "run a command in a spec's install environment, or dump its environment to screen or file"
+    "run a command in a spec's build environment, or dump its environment to screen or file"
 )
 section = "build"
 level = "long"
@@ -14,5 +14,7 @@
 setup_parser = env_utility.setup_parser
 
 
+#TODO create a dropin/dive funciton that dev_build and build_env cmd's can use
+# create a default shell utility for picking the current shell to start with
 def build_env(parser, args):
     env_utility.emulate_env_utility("build-env", Context.BUILD, args)
diff --git a/lib/spack/spack/cmd/common/env_utility.py b/lib/spack/spack/cmd/common/env_utility.py
index e80a090fc06..57988bcc5a6 100644
--- a/lib/spack/spack/cmd/common/env_utility.py
+++ b/lib/spack/spack/cmd/common/env_utility.py
@@ -10,6 +10,7 @@
 import spack.cmd
 import spack.deptypes as dt
 import spack.error
+import spack.prompt
 import spack.spec
 import spack.store
 from spack import build_environment, traverse
@@ -26,7 +27,9 @@ 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(
+        "--dive", action="store_true", help="dive into the build-env in a subshell"
+    )
     subparser.add_argument(
         "spec",
         nargs=argparse.REMAINDER,
@@ -124,7 +127,7 @@ def emulate_env_utility(cmd_name, context: Context, args):
             ),
         )
 
-    build_environment.setup_package(spec.package, args.dirty, context)
+    mods = build_environment.setup_package(spec.package, args.dirty, context)
 
     if args.dump:
         # Dump a source-able environment to a text file.
@@ -136,13 +139,14 @@ 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)
 
+    if args.dive:
+        mods.extend(spack.prompt.prompt_modifications(f"{spec.name}-build-env", cmd[0]))
+        os.execvp(cmd[0], [cmd[0]])
+
     elif not bool(args.pickle or args.dump):
         # If no command or dump/pickle option then act like the "env" command
         # and print out env vars.
diff --git a/lib/spack/spack/prompt.py b/lib/spack/spack/prompt.py
index 98bf6b0f246..1b6e4f0a18c 100644
--- a/lib/spack/spack/prompt.py
+++ b/lib/spack/spack/prompt.py
@@ -4,6 +4,7 @@
 # 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
@@ -19,16 +20,16 @@ def prompt_modifications(prompt, shell, env=os.environ):
         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:
+        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','$ ')}")
+                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', '$ ')}")
+                mods.set("PS1", f"{zsh_color_prompt} {env.get('PS1', '$ ')}")
         else:
-            mods.set('PS1', f"{prompt} {env.get('PS1', '$ ')}")
+            mods.set("PS1", f"{prompt} {env.get('PS1', '$ ')}")
 
     return mods
 
@@ -70,4 +71,3 @@ def custom_prompt(prompt, shell):
             """
         ).lstrip("\n")
     return cmds
-