spack env track command (spack#41897)

This PR adds a sub-command to `spack env` (`track`) which allows users to add/link
anonymous environments into their installation as named environments. This allows
users to more easily track their installed packages and the environments they're
dependencies of. For example, with the addition of spack#41731 it's now easier to remove
all packages not required by any environments with,

```
spack gc -bE
```

#### Usage
```
spack env track /path/to/env
==> Linked environment in /path/to/env
==> You can activate this environment with:
==>     spack env activate env
```

By default `track /path/to/env` will use the last directory in the path as the name of 
the environment. However users may customize the name of the linked environment
with `-n | --name`. Shown below.
```
spack env track /path/to/env --name foo 
==> Tracking environment in /path/to/env
==> You can activate this environment with:
==>     spack env activate foo
```

When removing a linked environment, Spack will remove the link to the environment
but will keep the structure of the environment within the directory. This will allow
users to remove a linked environment from their installation without deleting it from
a shared repository.

There is a `spack env untrack` command that can be used to *only* untrack a tracked
environment -- it will fail if it is used on a managed environment.  Users can also use
`spack env remove` to untrack an environment.

This allows users to continue to share environments in git repositories  while also having
the dependencies of those environments be remembered by Spack.

---------

Co-authored-by: Todd Gamblin <tgamblin@llnl.gov>

lib/spack/spack/cmd/env.py

@@ -10,11 +10,12 @@
 import sys
 import tempfile
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Optional, Set
 
 import llnl.string as string
 import llnl.util.filesystem as fs
 import llnl.util.tty as tty
+from llnl.util.symlink import islink, symlink
 from llnl.util.tty.colify import colify
 from llnl.util.tty.color import cescape, colorize
 
@@ -50,6 +51,8 @@
     "update",
     "revert",
     "depfile",
+    "track",
+    "untrack",
 ]
 
 
@@ -447,78 +450,220 @@ def env_deactivate(args):
 
 
 #
-# env remove
+# env track
 #
-def env_remove_setup_parser(subparser):
-    """remove managed environment(s)
+def env_track_setup_parser(subparser):
+    """track an environment from a directory in Spack"""
+    subparser.add_argument("-n", "--name", help="custom environment name")
+    subparser.add_argument("dir", help="path to environment")
+    arguments.add_common_arguments(subparser, ["yes_to_all"])
 
-    remove existing environment(s) managed by Spack
 
-    directory environments and manifests embedded in repositories must be
-    removed manually
-    """
-    subparser.add_argument(
-        "rm_env", metavar="env", nargs="+", help="name(s) of the environment(s) being removed"
-    )
-    arguments.add_common_arguments(subparser, ["yes_to_all"])
-    subparser.add_argument(
-        "-f",
-        "--force",
-        action="store_true",
-        help="force removal even when included in other environment(s)",
+def env_track(args):
+    src_path = os.path.abspath(args.dir)
+    if not ev.is_env_dir(src_path):
+        tty.die("Cannot track environment. Path doesn't contain an environment")
+
+    if args.name:
+        name = args.name
+    else:
+        name = os.path.basename(src_path)
+
+    try:
+        dst_path = ev.environment_dir_from_name(name, exists_ok=False)
+    except ev.SpackEnvironmentError:
+        tty.die(
+            f"An environment named {name} already exists. Set a name with:"
+            "\n\n"
+            f"        spack env track --name NAME {src_path}\n"
+        )
+
+    symlink(src_path, dst_path)
+
+    tty.msg(f"Tracking environment in {src_path}")
+    tty.msg(
+        "You can now activate this environment with the following command:\n\n"
+        f"        spack env activate {name}\n"
     )
 
 
-def env_remove(args):
-    """remove existing environment(s)"""
-    remove_envs = []
-    valid_envs = []
-    bad_envs = []
+#
+# env remove & untrack helpers
+#
+def filter_managed_env_names(env_names: Set[str]) -> Set[str]:
+    tracked_env_names = {e for e in env_names if islink(ev.environment_dir_from_name(e))}
+    managed_env_names = env_names - set(tracked_env_names)
+
+    num_managed_envs = len(managed_env_names)
+    managed_envs_str = " ".join(managed_env_names)
+    if num_managed_envs >= 2:
+        tty.error(
+            f"The following are not tracked environments. "
+            "To remove them completely run,"
+            "\n\n"
+            f"        spack env rm {managed_envs_str}\n"
+        )
+
+    elif num_managed_envs > 0:
+        tty.error(
+            f"'{managed_envs_str}' is not a tracked env. "
+            "To remove it completely run,"
+            "\n\n"
+            f"        spack env rm {managed_envs_str}\n"
+        )
+
+    return tracked_env_names
+
 
-    for env_name in ev.all_environment_names():
+def get_valid_envs(env_names: Set[str]) -> Set[ev.Environment]:
+    valid_envs = set()
+    for env_name in env_names:
         try:
             env = ev.read(env_name)
-            valid_envs.append(env)
+            valid_envs.add(env)
 
-            if env_name in args.rm_env:
-                remove_envs.append(env)
         except (spack.config.ConfigFormatError, ev.SpackEnvironmentConfigError):
-            if env_name in args.rm_env:
-                bad_envs.append(env_name)
+            pass
 
-    # Check if remove_env is included from another env before trying to remove
-    for env in valid_envs:
-        for remove_env in remove_envs:
-            # don't check if environment is included to itself
+    return valid_envs
+
+
+def _env_untrack_or_remove(
+    env_names: List[str], remove: bool = False, force: bool = False, yes_to_all: bool = False
+):
+    all_env_names = set(ev.all_environment_names())
+    known_env_names = set(env_names).intersection(all_env_names)
+    unknown_env_names = set(env_names) - known_env_names
+
+    # print error for unknown environments
+    for env_name in unknown_env_names:
+        tty.error(f"Environment '{env_name}' does not exist")
+
+    # if only unlinking is allowed, remove all environments
+    # which do not point internally at symlinks
+    if not remove:
+        env_names_to_remove = filter_managed_env_names(known_env_names)
+    else:
+        env_names_to_remove = known_env_names
+
+    # initalize all environments with valid spack.yaml configs
+    all_valid_envs = get_valid_envs(all_env_names)
+
+    # build a task list of environments and bad env names to remove
+    envs_to_remove = [e for e in all_valid_envs if e.name in env_names_to_remove]
+    bad_env_names_to_remove = env_names_to_remove - {e.name for e in envs_to_remove}
+    for remove_env in envs_to_remove:
+        for env in all_valid_envs:
+            # don't check if an environment is included to itself
             if env.name == remove_env.name:
                 continue
 
+            # check if an environment is included un another
             if remove_env.path in env.included_concrete_envs:
-                msg = f'Environment "{remove_env.name}" is being used by environment "{env.name}"'
-                if args.force:
+                msg = f"Environment '{remove_env.name}' is used by environment '{env.name}'"
+                if force:
                     tty.warn(msg)
                 else:
-                    tty.die(msg)
-
-    if not args.yes_to_all:
-        environments = string.plural(len(args.rm_env), "environment", show_n=False)
-        envs = string.comma_and(args.rm_env)
-        answer = tty.get_yes_or_no(f"Really remove {environments} {envs}?", default=False)
+                    tty.error(msg)
+                    envs_to_remove.remove(remove_env)
+
+    # ask the user if they really want to remove the known environments
+    # force should do the same as yes to all here following the symantics of rm
+    if not (yes_to_all or force) and (envs_to_remove or bad_env_names_to_remove):
+        environments = string.plural(len(env_names_to_remove), "environment", show_n=False)
+        envs = string.comma_and(list(env_names_to_remove))
+        answer = tty.get_yes_or_no(
+            f"Really {'remove' if remove else 'untrack'} {environments} {envs}?", default=False
+        )
         if not answer:
             tty.die("Will not remove any environments")
 
-    for env in remove_envs:
+    # keep track of the environments we remove for later printing the exit code
+    removed_env_names = []
+    for env in envs_to_remove:
         name = env.name
-        if env.active:
-            tty.die(f"Environment {name} can't be removed while activated.")
-        env.destroy()
-        tty.msg(f"Successfully removed environment '{name}'")
+        if not force and env.active:
+            tty.error(
+                f"Environment '{name}' can't be "
+                f"{'removed' if remove else 'untracked'} while activated."
+            )
+            continue
+        # Get path to check if environment is a tracked / symlinked environment
+        if islink(env.path):
+            real_env_path = os.path.realpath(env.path)
+            os.unlink(env.path)
+            tty.msg(
+                f"Sucessfully untracked environment '{name}', "
+                "but it can still be found at:\n\n"
+                f"        {real_env_path}\n"
+            )
+        else:
+            env.destroy()
+            tty.msg(f"Successfully removed environment '{name}'")
+
+        removed_env_names.append(env.name)
 
-    for bad_env_name in bad_envs:
+    for bad_env_name in bad_env_names_to_remove:
         shutil.rmtree(
             spack.environment.environment.environment_dir_from_name(bad_env_name, exists_ok=True)
         )
         tty.msg(f"Successfully removed environment '{bad_env_name}'")
+        removed_env_names.append(env.name)
+
+    # Following the design of linux rm we should exit with a status of 1
+    # anytime we cannot delete every environment the user asks for.
+    # However, we should still process all the environments we know about
+    # and delete them instead of failing on the first unknown enviornment.
+    if len(removed_env_names) < len(known_env_names):
+        sys.exit(1)
+
+
+#
+# env untrack
+#
+def env_untrack_setup_parser(subparser):
+    """track an environment from a directory in Spack"""
+    subparser.add_argument("env", nargs="+", help="tracked environment name")
+    subparser.add_argument(
+        "-f", "--force", action="store_true", help="force unlink even when environment is active"
+    )
+    arguments.add_common_arguments(subparser, ["yes_to_all"])
+
+
+def env_untrack(args):
+    _env_untrack_or_remove(
+        env_names=args.env, force=args.force, yes_to_all=args.yes_to_all, remove=False
+    )
+
+
+#
+# env remove
+#
+def env_remove_setup_parser(subparser):
+    """remove managed environment(s)
+
+    remove existing environment(s) managed by Spack
+
+    directory environments and manifests embedded in repositories must be
+    removed manually
+    """
+    subparser.add_argument(
+        "rm_env", metavar="env", nargs="+", help="name(s) of the environment(s) being removed"
+    )
+    arguments.add_common_arguments(subparser, ["yes_to_all"])
+    subparser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        help="force removal even when included in other environment(s)",
+    )
+
+
+def env_remove(args):
+    """remove existing environment(s)"""
+    _env_untrack_or_remove(
+        env_names=args.rm_env, remove=True, force=args.force, yes_to_all=args.yes_to_all
+    )
 
 
 #

lib/spack/spack/environment/environment.py

@@ -20,7 +20,7 @@
 import llnl.util.tty as tty
 import llnl.util.tty.color as clr
 from llnl.util.link_tree import ConflictingSpecsError
-from llnl.util.symlink import readlink, symlink
+from llnl.util.symlink import islink, readlink, symlink
 
 import spack
 import spack.caches
@@ -668,7 +668,7 @@ def from_dict(base_path, d):
 
     @property
     def _current_root(self):
-        if not os.path.islink(self.root):
+        if not islink(self.root):
             return None
 
         root = readlink(self.root)