Support `spack env activate --with-view <name> <env>` (#40549)

Currently `spack env activate --with-view` exists, but is a no-op.

So, it is not too much of a breaking change to make this redundant flag
accept a value `spack env activate --with-view <name>` which activates
a particular view by name.

The view name is stored in `SPACK_ENV_VIEW`.

This also fixes an issue where deactivating a view that was activated
with `--without-view` possibly removes entries from PATH, since now we
keep track of whether the default view was "enabled" or not.

5 files changed, 114 insertions, 69 deletions

diff --git a/lib/spack/spack/cmd/env.py b/lib/spack/spack/cmd/env.py
index 6c22e70a5d..cf5671aafa 100644
--- a/lib/spack/spack/cmd/env.py
+++ b/lib/spack/spack/cmd/env.py
@@ -8,6 +8,7 @@ import os
 import shutil
 import sys
 import tempfile
+from typing import Optional
 
 import llnl.string as string
 import llnl.util.filesystem as fs
@@ -96,22 +97,16 @@ def env_activate_setup_parser(subparser):
 
     view_options = subparser.add_mutually_exclusive_group()
     view_options.add_argument(
-        "-v",
         "--with-view",
-        action="store_const",
-        dest="with_view",
-        const=True,
-        default=True,
-        help="update PATH, etc., with associated view",
+        "-v",
+        metavar="name",
+        help="set runtime environment variables for specific view",
     )
     view_options.add_argument(
-        "-V",
         "--without-view",
-        action="store_const",
-        dest="with_view",
-        const=False,
-        default=True,
-        help="do not update PATH, etc., with associated view",
+        "-V",
+        action="store_true",
+        help="do not set runtime environment variables for any view",
     )
 
     subparser.add_argument(
@@ -197,10 +192,20 @@ def env_activate(args):
 
     # Activate new environment
     active_env = ev.Environment(env_path)
+
+    # Check if runtime environment variables are requested, and if so, for what view.
+    view: Optional[str] = None
+    if args.with_view:
+        view = args.with_view
+        if not active_env.has_view(view):
+            tty.die(f"The environment does not have a view named '{view}'")
+    elif not args.without_view and active_env.has_view(ev.default_view_name):
+        view = ev.default_view_name
+
     cmds += spack.environment.shell.activate_header(
-        env=active_env, shell=args.shell, prompt=env_prompt if args.prompt else None
+        env=active_env, shell=args.shell, prompt=env_prompt if args.prompt else None, view=view
     )
-    env_mods.extend(spack.environment.shell.activate(env=active_env, add_view=args.with_view))
+    env_mods.extend(spack.environment.shell.activate(env=active_env, view=view))
     cmds += env_mods.shell_modifications(args.shell)
     sys.stdout.write(cmds)
 
diff --git a/lib/spack/spack/environment/__init__.py b/lib/spack/spack/environment/__init__.py
index 227b48670c..ac598e8421 100644
--- a/lib/spack/spack/environment/__init__.py
+++ b/lib/spack/spack/environment/__init__.py
@@ -365,6 +365,7 @@ from .environment import (
     read,
     root,
     spack_env_var,
+    spack_env_view_var,
     update_yaml,
 )
 
@@ -397,5 +398,6 @@ __all__ = [
     "read",
     "root",
     "spack_env_var",
+    "spack_env_view_var",
     "update_yaml",
 ]
diff --git a/lib/spack/spack/environment/environment.py b/lib/spack/spack/environment/environment.py
index 496a8b332a..ee48955ac5 100644
--- a/lib/spack/spack/environment/environment.py
+++ b/lib/spack/spack/environment/environment.py
@@ -64,6 +64,8 @@ from spack.variant import UnknownVariantError
 #: environment variable used to indicate the active environment
 spack_env_var = "SPACK_ENV"
 
+#: environment variable used to indicate the active environment view
+spack_env_view_var = "SPACK_ENV_VIEW"
 
 #: currently activated environment
 _active_environment: Optional["Environment"] = None
@@ -1595,16 +1597,14 @@ class Environment:
 
     @property
     def default_view(self):
-        if not self.views:
-            raise SpackEnvironmentError("{0} does not have a view enabled".format(self.name))
-
-        if default_view_name not in self.views:
-            raise SpackEnvironmentError(
-                "{0} does not have a default view enabled".format(self.name)
-            )
+        if not self.has_view(default_view_name):
+            raise SpackEnvironmentError(f"{self.name} does not have a default view enabled")
 
         return self.views[default_view_name]
 
+    def has_view(self, view_name: str) -> bool:
+        return view_name in self.views
+
     def update_default_view(self, path_or_bool: Union[str, bool]) -> None:
         """Updates the path of the default view.
 
@@ -1690,14 +1690,14 @@ class Environment:
                 "Loading the environment view will require reconcretization." % self.name
             )
 
-    def _env_modifications_for_default_view(self, reverse=False):
+    def _env_modifications_for_view(self, view: ViewDescriptor, reverse: bool = False):
         all_mods = spack.util.environment.EnvironmentModifications()
 
         visited = set()
 
         errors = []
         for root_spec in self.concrete_roots():
-            if root_spec in self.default_view and root_spec.installed and root_spec.package:
+            if root_spec in view and root_spec.installed and root_spec.package:
                 for spec in root_spec.traverse(deptype="run", root=True):
                     if spec.name in visited:
                         # It is expected that only one instance of the package
@@ -1714,7 +1714,7 @@ class Environment:
                         visited.add(spec.name)
 
                     try:
-                        mods = uenv.environment_modifications_for_spec(spec, self.default_view)
+                        mods = uenv.environment_modifications_for_spec(spec, view)
                     except Exception as e:
                         msg = "couldn't get environment settings for %s" % spec.format(
                             "{name}@{version} /{hash:7}"
@@ -1726,22 +1726,22 @@ class Environment:
 
         return all_mods, errors
 
-    def add_default_view_to_env(self, env_mod):
-        """
-        Collect the environment modifications to activate an environment using the
-        default view. Removes duplicate paths.
+    def add_view_to_env(
+        self, env_mod: spack.util.environment.EnvironmentModifications, view: str
+    ) -> spack.util.environment.EnvironmentModifications:
+        """Collect the environment modifications to activate an environment using the provided
+        view. Removes duplicate paths.
 
         Args:
-            env_mod (spack.util.environment.EnvironmentModifications): the environment
-                modifications object that is modified.
-        """
-        if default_view_name not in self.views:
-            # No default view to add to shell
+            env_mod: the environment modifications object that is modified.
+            view: the name of the view to activate."""
+        descriptor = self.views.get(view)
+        if not descriptor:
             return env_mod
 
-        env_mod.extend(uenv.unconditional_environment_modifications(self.default_view))
+        env_mod.extend(uenv.unconditional_environment_modifications(descriptor))
 
-        mods, errors = self._env_modifications_for_default_view()
+        mods, errors = self._env_modifications_for_view(descriptor)
         env_mod.extend(mods)
         if errors:
             for err in errors:
@@ -1753,22 +1753,22 @@ class Environment:
 
         return env_mod
 
-    def rm_default_view_from_env(self, env_mod):
-        """
-        Collect the environment modifications to deactivate an environment using the
-        default view. Reverses the action of ``add_default_view_to_env``.
+    def rm_view_from_env(
+        self, env_mod: spack.util.environment.EnvironmentModifications, view: str
+    ) -> spack.util.environment.EnvironmentModifications:
+        """Collect the environment modifications to deactivate an environment using the provided
+        view. Reverses the action of ``add_view_to_env``.
 
         Args:
-            env_mod (spack.util.environment.EnvironmentModifications): the environment
-                modifications object that is modified.
-        """
-        if default_view_name not in self.views:
-            # No default view to add to shell
+            env_mod: the environment modifications object that is modified.
+            view: the name of the view to deactivate."""
+        descriptor = self.views.get(view)
+        if not descriptor:
             return env_mod
 
-        env_mod.extend(uenv.unconditional_environment_modifications(self.default_view).reversed())
+        env_mod.extend(uenv.unconditional_environment_modifications(descriptor).reversed())
 
-        mods, _ = self._env_modifications_for_default_view(reverse=True)
+        mods, _ = self._env_modifications_for_view(descriptor, reverse=True)
         env_mod.extend(mods)
 
         return env_mod
diff --git a/lib/spack/spack/environment/shell.py b/lib/spack/spack/environment/shell.py
index 380e49fa0f..a4f9634a8d 100644
--- a/lib/spack/spack/environment/shell.py
+++ b/lib/spack/spack/environment/shell.py
@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 import os
+from typing import Optional
 
 import llnl.util.tty as tty
 from llnl.util.tty.color import colorize
@@ -13,12 +14,14 @@ import spack.store
 from spack.util.environment import EnvironmentModifications
 
 
-def activate_header(env, shell, prompt=None):
+def activate_header(env, shell, prompt=None, view: Optional[str] = None):
     # Construct the commands to run
     cmds = ""
     if shell == "csh":
         # TODO: figure out how to make color work for csh
         cmds += "setenv SPACK_ENV %s;\n" % env.path
+        if view:
+            cmds += "setenv SPACK_ENV_VIEW %s;\n" % view
         cmds += 'alias despacktivate "spack env deactivate";\n'
         if prompt:
             cmds += "if (! $?SPACK_OLD_PROMPT ) "
@@ -29,6 +32,8 @@ def activate_header(env, shell, prompt=None):
             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"
@@ -40,15 +45,21 @@ def activate_header(env, shell, prompt=None):
     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:
         if "color" in os.getenv("TERM", "") and prompt:
             prompt = colorize("@G{%s}" % prompt, color=True, enclose=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 += "if [ -z ${SPACK_OLD_PS1+x} ]; then\n"
@@ -66,12 +77,14 @@ def deactivate_header(shell):
     cmds = ""
     if shell == "csh":
         cmds += "unsetenv SPACK_ENV;\n"
+        cmds += "unsetenv SPACK_ENV_VIEW;\n"
         cmds += "if ( $?SPACK_OLD_PROMPT ) "
         cmds += '    eval \'set prompt="$SPACK_OLD_PROMPT" &&'
         cmds += "          unsetenv SPACK_OLD_PROMPT';\n"
         cmds += "unalias despacktivate;\n"
     elif shell == "fish":
         cmds += "set -e SPACK_ENV;\n"
+        cmds += "set -e SPACK_ENV_VIEW;\n"
         cmds += "functions -e despacktivate;\n"
         #
         # NOTE: Not changing fish_prompt (above) => no need to restore it here.
@@ -79,14 +92,19 @@ def deactivate_header(shell):
     elif shell == "bat":
         # TODO: Color
         cmds += 'set "SPACK_ENV="\n'
+        cmds += 'set "SPACK_ENV_VIEW="\n'
         # TODO: despacktivate
         # TODO: prompt
     elif shell == "pwsh":
         cmds += "Set-Item -Path Env:SPACK_ENV\n"
+        cmds += "Set-Item -Path Env:SPACK_ENV_VIEW\n"
     else:
         cmds += "if [ ! -z ${SPACK_ENV+x} ]; then\n"
         cmds += "unset SPACK_ENV; export SPACK_ENV;\n"
         cmds += "fi;\n"
+        cmds += "if [ ! -z ${SPACK_ENV_VIEW+x} ]; then\n"
+        cmds += "unset SPACK_ENV_VIEW; export SPACK_ENV_VIEW;\n"
+        cmds += "fi;\n"
         cmds += "alias despacktivate > /dev/null 2>&1 && unalias despacktivate;\n"
         cmds += "if [ ! -z ${SPACK_OLD_PS1+x} ]; then\n"
         cmds += "    if [ \"$SPACK_OLD_PS1\" = '$$$$' ]; then\n"
@@ -100,24 +118,23 @@ def deactivate_header(shell):
     return cmds
 
 
-def activate(env, use_env_repo=False, add_view=True):
-    """
-    Activate an environment and append environment modifications
+def activate(
+    env: ev.Environment, use_env_repo=False, view: Optional[str] = "default"
+) -> EnvironmentModifications:
+    """Activate an environment and append environment modifications
 
     To activate an environment, we add its configuration scope to the
     existing Spack configuration, and we set active to the current
     environment.
 
     Arguments:
-        env (spack.environment.Environment): the environment to activate
-        use_env_repo (bool): use the packages exactly as they appear in the
-            environment's repository
-        add_view (bool): generate commands to add view to path variables
+        env: the environment to activate
+        use_env_repo: use the packages exactly as they appear in the environment's repository
+        view: generate commands to add runtime environment variables for named view
 
     Returns:
         spack.util.environment.EnvironmentModifications: Environment variables
-        modifications to activate environment.
-    """
+        modifications to activate environment."""
     ev.activate(env, use_env_repo=use_env_repo)
 
     env_mods = EnvironmentModifications()
@@ -129,9 +146,9 @@ def activate(env, use_env_repo=False, add_view=True):
     # become PATH variables.
     #
     try:
-        if add_view and ev.default_view_name in env.views:
+        if view and env.has_view(view):
             with spack.store.STORE.db.read_transaction():
-                env.add_default_view_to_env(env_mods)
+                env.add_view_to_env(env_mods, view)
     except (spack.repo.UnknownPackageError, spack.repo.UnknownNamespaceError) as e:
         tty.error(e)
         tty.die(
@@ -145,17 +162,15 @@ def activate(env, use_env_repo=False, add_view=True):
     return env_mods
 
 
-def deactivate():
-    """
-    Deactivate an environment and collect corresponding environment modifications.
+def deactivate() -> EnvironmentModifications:
+    """Deactivate an environment and collect corresponding environment modifications.
 
     Note: unloads the environment in its current state, not in the state it was
         loaded in, meaning that specs that were removed from the spack environment
         after activation are not unloaded.
 
     Returns:
-        spack.util.environment.EnvironmentModifications: Environment variables
-        modifications to activate environment.
+        Environment variables modifications to activate environment.
     """
     env_mods = EnvironmentModifications()
     active = ev.active_environment()
@@ -163,10 +178,12 @@ def deactivate():
     if active is None:
         return env_mods
 
-    if ev.default_view_name in active.views:
+    active_view = os.getenv(ev.spack_env_view_var)
+
+    if active_view and active.has_view(active_view):
         try:
             with spack.store.STORE.db.read_transaction():
-                active.rm_default_view_from_env(env_mods)
+                active.rm_view_from_env(env_mods, active_view)
         except (spack.repo.UnknownPackageError, spack.repo.UnknownNamespaceError) as e:
             tty.warn(e)
             tty.warn(
diff --git a/lib/spack/spack/test/cmd/env.py b/lib/spack/spack/test/cmd/env.py
index cef5ccbcd5..4845d12206 100644
--- a/lib/spack/spack/test/cmd/env.py
+++ b/lib/spack/spack/test/cmd/env.py
@@ -663,7 +663,7 @@ packages:
         e.write()
 
         env_mod = spack.util.environment.EnvironmentModifications()
-        e.add_default_view_to_env(env_mod)
+        e.add_view_to_env(env_mod, "default")
         env_variables = {}
         env_mod.apply_modifications(env_variables)
         assert str(fake_bin) in env_variables["PATH"]
@@ -2356,7 +2356,7 @@ def test_env_activate_sh_prints_shell_output(tmpdir, mock_stage, mock_fetch, ins
     This is a cursory check; ``share/spack/qa/setup-env-test.sh`` checks
     for correctness.
     """
-    env("create", "test", add_view=True)
+    env("create", "test")
 
     out = env("activate", "--sh", "test")
     assert "export SPACK_ENV=" in out
@@ -2371,7 +2371,7 @@ def test_env_activate_sh_prints_shell_output(tmpdir, mock_stage, mock_fetch, ins
 
 def test_env_activate_csh_prints_shell_output(tmpdir, mock_stage, mock_fetch, install_mockery):
     """Check the shell commands output by ``spack env activate --csh``."""
-    env("create", "test", add_view=True)
+    env("create", "test")
 
     out = env("activate", "--csh", "test")
     assert "setenv SPACK_ENV" in out
@@ -2388,7 +2388,7 @@ def test_env_activate_csh_prints_shell_output(tmpdir, mock_stage, mock_fetch, in
 def test_env_activate_default_view_root_unconditional(mutable_mock_env_path):
     """Check that the root of the default view in the environment is added
     to the shell unconditionally."""
-    env("create", "test", add_view=True)
+    env("create", "test")
 
     with ev.read("test") as e:
         viewdir = e.default_view.root
@@ -2403,6 +2403,27 @@ def test_env_activate_default_view_root_unconditional(mutable_mock_env_path):
     )
 
 
+def test_env_activate_custom_view(tmp_path: pathlib.Path, mock_packages):
+    """Check that an environment can be activated with a non-default view."""
+    env_template = tmp_path / "spack.yaml"
+    default_dir = tmp_path / "defaultdir"
+    nondefaultdir = tmp_path / "nondefaultdir"
+    with open(env_template, "w") as f:
+        f.write(
+            f"""\
+spack:
+  specs: [a]
+  view:
+    default:
+      root: {default_dir}
+    nondefault:
+      root: {nondefaultdir}"""
+        )
+    env("create", "test", str(env_template))
+    shell = env("activate", "--sh", "--with-view", "nondefault", "test")
+    assert os.path.join(nondefaultdir, "bin") in shell
+
+
 def test_concretize_user_specs_together():
     e = ev.create("coconcretization")
     e.unify = True