Improve `spack find` output in environments (#42334)

This adds some improvements to `spack find` output when in environments based
around some thoughts about what users want to know when they're in an env.

If you're working in an enviroment, you mostly care about:
* What are the roots
* Which ones are installed / not installed
* What's been added that still needs to be concretized

So, this PR adds a couple tweaks to display that information more clearly:

- [x] We now display install status next to every root. You can easily see
      which are installed and which aren't.

- [x] When you run `spack find -l` in an env, the roots now show their concrete
      hash (if they've been concretized). They previously would show `-------`
      (b/c the root spec itself is abstract), but showing the concretized root's
      hash is a lot more useful.

- [x] Newly added/unconcretized specs still show `-------`, which now makes more
      sense, b/c they are not concretized.

- [x] There is a new option, `-r` / `--only-roots` to *only* show env roots if
      you don't want to look at all the installed specs.

- [x] Roots in the installed spec list are now highlighted as bold. This is
      actually an old feature from the first env implementation , but various
      refactors had disabled it inadvertently.

2 files changed, 64 insertions, 38 deletions

diff --git a/lib/spack/spack/cmd/__init__.py b/lib/spack/spack/cmd/__init__.py
index b30c6360d9..bd6a3eb768 100644
--- a/lib/spack/spack/cmd/__init__.py
+++ b/lib/spack/spack/cmd/__init__.py
@@ -334,8 +334,7 @@ def display_specs(specs, args=None, **kwargs):
         variants (bool): Show variants with specs
         indent (int): indent each line this much
         groups (bool): display specs grouped by arch/compiler (default True)
-        decorators (dict): dictionary mappng specs to decorators
-        header_callback (typing.Callable): called at start of arch/compiler groups
+        decorator (typing.Callable): function to call to decorate specs
         all_headers (bool): show headers even when arch/compiler aren't defined
         output (typing.IO): A file object to write to. Default is ``sys.stdout``
 
@@ -384,15 +383,13 @@ def display_specs(specs, args=None, **kwargs):
         vfmt = "{variants}" if variants else ""
         format_string = nfmt + "{@version}" + ffmt + vfmt
 
-    transform = {"package": decorator, "fullpackage": decorator}
-
     def fmt(s, depth=0):
         """Formatter function for all output specs"""
         string = ""
         if hashes:
             string += gray_hash(s, hlen) + " "
         string += depth * "    "
-        string += s.cformat(format_string, transform=transform)
+        string += decorator(s, s.cformat(format_string))
         return string
 
     def format_list(specs):
@@ -451,7 +448,7 @@ def filter_loaded_specs(specs):
     return [x for x in specs if x.dag_hash() in hashes]
 
 
-def print_how_many_pkgs(specs, pkg_type=""):
+def print_how_many_pkgs(specs, pkg_type="", suffix=""):
     """Given a list of specs, this will print a message about how many
     specs are in that list.
 
@@ -462,7 +459,7 @@ def print_how_many_pkgs(specs, pkg_type=""):
             category, e.g. if pkg_type is "installed" then the message
             would be "3 installed packages"
     """
-    tty.msg("%s" % llnl.string.plural(len(specs), pkg_type + " package"))
+    tty.msg("%s" % llnl.string.plural(len(specs), pkg_type + " package") + suffix)
 
 
 def spack_is_git_repo():
diff --git a/lib/spack/spack/cmd/find.py b/lib/spack/spack/cmd/find.py
index f8e74f262f..d1917a73b5 100644
--- a/lib/spack/spack/cmd/find.py
+++ b/lib/spack/spack/cmd/find.py
@@ -3,7 +3,6 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
-import copy
 import sys
 
 import llnl.util.lang
@@ -14,6 +13,7 @@ import spack.bootstrap
 import spack.cmd as cmd
 import spack.environment as ev
 import spack.repo
+import spack.store
 from spack.cmd.common import arguments
 from spack.database import InstallStatuses
 
@@ -70,6 +70,12 @@ def setup_parser(subparser):
     arguments.add_common_arguments(subparser, ["long", "very_long", "tags", "namespaces"])
 
     subparser.add_argument(
+        "-r",
+        "--only-roots",
+        action="store_true",
+        help="don't show full list of installed specs in an environment",
+    )
+    subparser.add_argument(
         "-c",
         "--show-concretized",
         action="store_true",
@@ -189,26 +195,22 @@ def query_arguments(args):
     return q_args
 
 
-def setup_env(env):
+def make_env_decorator(env):
     """Create a function for decorating specs when in an environment."""
 
-    def strip_build(seq):
-        return set(s.copy(deps=("link", "run")) for s in seq)
-
-    added = set(strip_build(env.added_specs()))
-    roots = set(strip_build(env.roots()))
-    removed = set(strip_build(env.removed_specs()))
+    roots = set(env.roots())
+    removed = set(env.removed_specs())
 
     def decorator(spec, fmt):
         # add +/-/* to show added/removed/root specs
         if any(spec.dag_hash() == r.dag_hash() for r in roots):
-            return color.colorize("@*{%s}" % fmt)
+            return color.colorize(f"@*{{{fmt}}}")
         elif spec in removed:
-            return color.colorize("@K{%s}" % fmt)
+            return color.colorize(f"@K{{{fmt}}}")
         else:
-            return "%s" % fmt
+            return fmt
 
-    return decorator, added, roots, removed
+    return decorator
 
 
 def display_env(env, args, decorator, results):
@@ -223,28 +225,51 @@ def display_env(env, args, decorator, results):
     """
     tty.msg("In environment %s" % env.name)
 
-    if not env.user_specs:
-        tty.msg("No root specs")
-    else:
-        tty.msg("Root specs")
+    num_roots = len(env.user_specs) or "No"
+    tty.msg(f"{num_roots} root specs")
+
+    concrete_specs = {
+        root: concrete_root
+        for root, concrete_root in zip(env.concretized_user_specs, env.concrete_roots())
+    }
 
-        # Root specs cannot be displayed with prefixes, since those are not
-        # set for abstract specs. Same for hashes
-        root_args = copy.copy(args)
-        root_args.paths = False
+    def root_decorator(spec, string):
+        """Decorate root specs with their install status if needed"""
+        concrete = concrete_specs.get(spec)
+        if concrete:
+            status = color.colorize(concrete.install_status().value)
+            hash = concrete.dag_hash()
+        else:
+            status = color.colorize(spack.spec.InstallStatus.absent.value)
+            hash = "-" * 32
 
-        # Roots are displayed with variants, etc. so that we can see
-        # specifically what the user asked for.
+        # TODO: status has two extra spaces on the end of it, but fixing this and other spec
+        # TODO: space format idiosyncrasies is complicated. Fix this eventually
+        status = status[:-2]
+
+        if args.long or args.very_long:
+            hash = color.colorize(f"@K{{{hash[: 7 if args.long else None]}}}")
+            return f"{status} {hash} {string}"
+        else:
+            return f"{status} {string}"
+
+    with spack.store.STORE.db.read_transaction():
         cmd.display_specs(
             env.user_specs,
-            root_args,
-            decorator=lambda s, f: color.colorize("@*{%s}" % f),
+            args,
+            # these are overrides of CLI args
+            paths=False,
+            long=False,
+            very_long=False,
+            # these enforce details in the root specs to show what the user asked for
             namespaces=True,
             show_flags=True,
             show_full_compiler=True,
+            decorator=root_decorator,
             variants=True,
         )
-        print()
+
+    print()
 
     if args.show_concretized:
         tty.msg("Concretized roots")
@@ -254,7 +279,7 @@ def display_env(env, args, decorator, results):
     # Display a header for the installed packages section IF there are installed
     # packages. If there aren't any, we'll just end up printing "0 installed packages"
     # later.
-    if results:
+    if results and not args.only_roots:
         tty.msg("Installed packages")
 
 
@@ -263,9 +288,10 @@ def find(parser, args):
     results = args.specs(**q_args)
 
     env = ev.active_environment()
-    decorator = lambda s, f: f
-    if env:
-        decorator, _, roots, _ = setup_env(env)
+    if not env and args.only_roots:
+        tty.die("-r / --only-roots requires an active environment")
+
+    decorator = make_env_decorator(env) if env else lambda s, f: f
 
     # use groups by default except with format.
     if args.groups is None:
@@ -292,9 +318,12 @@ def find(parser, args):
             if env:
                 display_env(env, args, decorator, results)
 
-        cmd.display_specs(results, args, decorator=decorator, all_headers=True)
+        count_suffix = " (not shown)"
+        if not args.only_roots:
+            cmd.display_specs(results, args, decorator=decorator, all_headers=True)
+            count_suffix = ""
 
         # print number of installed packages last (as the list may be long)
         if sys.stdout.isatty() and args.groups:
             pkg_type = "loaded" if args.loaded else "installed"
-            spack.cmd.print_how_many_pkgs(results, pkg_type)
+            spack.cmd.print_how_many_pkgs(results, pkg_type, suffix=count_suffix)