From 9b3c200a071853189543c845777b7b8dc9dff6d5 Mon Sep 17 00:00:00 2001 From: psakiev Date: Mon, 11 Nov 2024 16:44:48 -0700 Subject: [PATCH] Add prompt capability to dev_build drop-in --- lib/spack/spack/build_environment.py | 4 +- lib/spack/spack/cmd/common/env_utility.py | 4 ++ lib/spack/spack/cmd/dev_build.py | 12 +++- lib/spack/spack/environment/shell.py | 45 ++------------ lib/spack/spack/prompt.py | 73 +++++++++++++++++++++++ 5 files changed, 96 insertions(+), 42 deletions(-) create mode 100644 lib/spack/spack/prompt.py diff --git a/lib/spack/spack/build_environment.py b/lib/spack/spack/build_environment.py index 73f5d438a26..bf0b2f5b0be 100644 --- a/lib/spack/spack/build_environment.py +++ b/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() diff --git a/lib/spack/spack/cmd/common/env_utility.py b/lib/spack/spack/cmd/common/env_utility.py index 6371ef65a83..e80a090fc06 100644 --- a/lib/spack/spack/cmd/common/env_utility.py +++ b/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) diff --git a/lib/spack/spack/cmd/dev_build.py b/lib/spack/spack/cmd/dev_build.py index a758c5b74f1..4eec718e902 100644 --- a/lib/spack/spack/cmd/dev_build.py +++ b/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]) diff --git a/lib/spack/spack/environment/shell.py b/lib/spack/spack/environment/shell.py index b1d87a48fd7..6856046f7de 100644 --- a/lib/spack/spack/environment/shell.py +++ b/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( - 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") + + if prompt: + cmds += spack.prompt.custom_prompt(prompt, shell) + return cmds diff --git a/lib/spack/spack/prompt.py b/lib/spack/spack/prompt.py new file mode 100644 index 00000000000..98bf6b0f246 --- /dev/null +++ b/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 +