refactor: clean up `spack find`, make `spack find -dp` work properly

Refactor `spack.cmd.display_specs()` and `spack find` so that any options
can be used together with -d. This cleans up the display logic
considerably, as there are no longer multiple "modes".

4 files changed, 132 insertions, 153 deletions

diff --git a/lib/spack/spack/cmd/__init__.py b/lib/spack/spack/cmd/__init__.py
index 4c80025e05..f47a4602ce 100644
--- a/lib/spack/spack/cmd/__init__.py
+++ b/lib/spack/spack/cmd/__init__.py
@@ -201,38 +201,55 @@ def gray_hash(spec, length):
     return colorize('@K{%s}' % h)
 
 
-def display_formatted_specs(specs, format_string, deps=False):
-    """Print a list of specs formatted with the provided string.
-
-    Arguments:
-        specs (list): list of specs to display.
-        deps (bool): whether to also print dependencies of specs.
-    """
-    for spec in specs:
-        print(spec.format(format_string))
-        if deps:
-            for depth, dep in spec.traverse(depth=True, root=False):
-                print("   " * depth, dep.format(format_string))
-
-
 def display_specs_as_json(specs, deps=False):
     """Convert specs to a list of json records."""
     seen = set()
     records = []
     for spec in specs:
-        if spec.dag_hash() not in seen:
-            seen.add(spec.dag_hash())
-            records.append(spec.to_record_dict())
+        if spec.dag_hash() in seen:
+            continue
+        seen.add(spec.dag_hash())
+        records.append(spec.to_record_dict())
 
         if deps:
             for dep in spec.traverse():
-                if dep.dag_hash() not in seen:
-                    seen.add(spec.dag_hash())
-                    records.append(dep.to_record_dict())
+                if dep.dag_hash() in seen:
+                    continue
+                seen.add(dep.dag_hash())
+                records.append(dep.to_record_dict())
 
     sjson.dump(records, sys.stdout)
 
 
+def iter_sections(specs, indent, all_headers):
+    """Break a list of specs into sections indexed by arch/compiler."""
+    # Make a dict with specs keyed by architecture and compiler.
+    index = index_by(specs, ('architecture', 'compiler'))
+    ispace = indent * ' '
+
+    # Traverse the index and print out each package
+    for i, (architecture, compiler) in enumerate(sorted(index)):
+        if i > 0:
+            print()
+
+        header = "%s{%s} / %s{%s}" % (
+            spack.spec.architecture_color,
+            architecture if architecture else 'no arch',
+            spack.spec.compiler_color,
+            compiler if compiler else 'no compiler')
+
+        # Sometimes we want to display specs that are not yet concretized.
+        # If they don't have a compiler / architecture attached to them,
+        # then skip the header
+        if all_headers or (architecture is not None or compiler is not None):
+            sys.stdout.write(ispace)
+            tty.hline(colorize(header), char='-')
+
+        specs = index[(architecture, compiler)]
+        specs.sort()
+        yield specs
+
+
 def display_specs(specs, args=None, **kwargs):
     """Display human readable specs with customizable formatting.
 
@@ -240,11 +257,7 @@ def display_specs(specs, args=None, **kwargs):
     arguments provided.
 
     Specs are grouped by architecture and compiler, and columnized if
-    possible.  There are three possible "modes":
-
-      * ``short`` (default): short specs with name and version, columnized
-      * ``paths``: Two columns: one for specs, one for paths
-      * ``deps``: Dependency-tree style, like ``spack spec``; can get long
+    possible.
 
     Options can add more information to the default display. Options can
     be provided either as keyword arguments or as an argparse namespace.
@@ -257,16 +270,19 @@ def display_specs(specs, args=None, **kwargs):
             formatting arguments
 
     Keyword Args:
-        mode (str): Either 'short', 'paths', or 'deps'
+        paths (bool): Show paths with each displayed spec
+        deps (bool): Display dependencies with specs
         long (bool): Display short hashes with specs
         very_long (bool): Display full hashes with specs (supersedes ``long``)
         namespace (bool): Print namespaces along with names
         show_flags (bool): Show compiler flags with specs
         variants (bool): Show variants with specs
         indent (int): indent each line this much
+        sections (bool): display specs grouped by arch/compiler (default True)
         decorators (dict): dictionary mappng specs to decorators
         header_callback (function): called at start of arch/compiler sections
         all_headers (bool): show headers even when arch/compiler aren't defined
+
     """
     def get_arg(name, default=None):
         """Prefer kwargs, then args, then default."""
@@ -277,12 +293,14 @@ def display_specs(specs, args=None, **kwargs):
         else:
             return default
 
-    mode          = get_arg('mode', 'short')
+    paths         = get_arg('paths', False)
+    deps          = get_arg('deps', False)
     hashes        = get_arg('long', False)
     namespace     = get_arg('namespace', False)
     flags         = get_arg('show_flags', False)
     full_compiler = get_arg('show_full_compiler', False)
     variants      = get_arg('variants', False)
+    sections      = get_arg('sections', True)
     all_headers   = get_arg('all_headers', False)
 
     decorator     = get_arg('decorator', None)
@@ -290,98 +308,70 @@ def display_specs(specs, args=None, **kwargs):
         decorator = lambda s, f: f
 
     indent = get_arg('indent', 0)
-    ispace = indent * ' '
 
     hlen = 7
     if get_arg('very_long', False):
         hashes = True
         hlen = None
 
-    nfmt = '{namespace}.{name}' if namespace else '{name}'
-    ffmt = ''
-    if full_compiler or flags:
-        ffmt += '{%compiler.name}'
-        if full_compiler:
-            ffmt += '{@compiler.version}'
-        ffmt += ' {compiler_flags}'
-    vfmt = '{variants}' if variants else ''
-    format_string = nfmt + '{@version}' + ffmt + vfmt
+    format_string = get_arg('format', None)
+    if format_string is None:
+        nfmt = '{namespace}.{name}' if namespace else '{name}'
+        ffmt = ''
+        if full_compiler or flags:
+            ffmt += '{%compiler.name}'
+            if full_compiler:
+                ffmt += '{@compiler.version}'
+            ffmt += ' {compiler_flags}'
+        vfmt = '{variants}' if variants else ''
+        format_string = nfmt + '{@version}' + ffmt + vfmt
 
-    # Make a dict with specs keyed by architecture and compiler.
-    index = index_by(specs, ('architecture', 'compiler'))
     transform = {'package': decorator, 'fullpackage': decorator}
 
-    # Traverse the index and print out each package
-    for i, (architecture, compiler) in enumerate(sorted(index)):
-        if i > 0:
-            print()
-
-        header = "%s{%s} / %s{%s}" % (
-            spack.spec.architecture_color,
-            architecture if architecture else 'no arch',
-            spack.spec.compiler_color,
-            compiler if compiler else 'no compiler')
-
-        # Sometimes we want to display specs that are not yet concretized.
-        # If they don't have a compiler / architecture attached to them,
-        # then skip the header
-        if all_headers or (architecture is not None or compiler is not None):
-            sys.stdout.write(ispace)
-            tty.hline(colorize(header), char='-')
-
-        specs = index[(architecture, compiler)]
-        specs.sort()
-
-        if mode == 'paths':
-            # Print one spec per line along with prefix path
-            abbreviated = [s.cformat(format_string, transform=transform)
-                           for s in specs]
-            width = max(len(s) for s in abbreviated)
-            width += 2
-
-            for abbrv, spec in zip(abbreviated, specs):
-                # optional hash prefix for paths
-                h = gray_hash(spec, hlen) if hashes else ''
-
-                # only show prefix for concrete specs
-                prefix = spec.prefix if spec.concrete else ''
-
-                # print it all out at once
-                fmt = "%%s%%s    %%-%ds%%s" % width
-                print(fmt % (ispace, h, abbrv, prefix))
-
-        elif mode == 'deps':
-            for spec in specs:
-                print(spec.tree(
-                    format=format_string,
-                    indent=4,
-                    prefix=(lambda s: gray_hash(s, hlen)) if hashes else None))
-
-        elif mode == 'short':
-            def fmt(s):
-                string = ""
-                if hashes:
-                    string += gray_hash(s, hlen) + ' '
-                string += s.cformat(
-                    nfmt + '{@version}' + vfmt, transform=transform)
-                return string
-
-            if not flags and not full_compiler:
-                # Print columns of output if not printing flags
-                colify((fmt(s) for s in specs), indent=indent)
-
+    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)
+        return string
+
+    def format_list(specs):
+        """Display a single list of specs, with no sections"""
+        # create the final, formatted versions of all specs
+        formatted = []
+        for spec in specs:
+            formatted.append((fmt(spec), spec))
+            if deps:
+                for depth, dep in spec.traverse(root=False, depth=True):
+                    formatted.append((fmt(dep, depth), dep))
+                formatted.append(('', None))  # mark newlines
+
+        # unless any of these are set, we can just colify and be done.
+        if not any((deps, paths)):
+            colify((f[0] for f in formatted), indent=indent)
+            return
+
+        # otherwise, we'll print specs one by one
+        max_width = max(len(f[0]) for f in formatted)
+        path_fmt = "%%-%ds%%s" % (max_width + 2)
+
+        for string, spec in formatted:
+            if not string:
+                print()  # print newline from above
+                continue
+
+            if paths:
+                print(path_fmt % (string, spec.prefix))
             else:
-                # Print one entry per line if including flags
-                for spec in specs:
-                    # Print the hash if necessary
-                    hsh = gray_hash(spec, hlen) + ' ' if hashes else ''
-                    print(ispace + hsh + spec.cformat(
-                        format_string, transform=transform))
+                print(string)
 
-        else:
-            raise ValueError(
-                "Invalid mode for display_specs: %s. Must be one of (paths,"
-                "deps, short)." % mode)
+    if sections:
+        for specs in iter_sections(specs, indent, all_headers):
+            format_list(specs)
+    else:
+        format_list(sorted(specs))
 
 
 def spack_is_git_repo():
diff --git a/lib/spack/spack/cmd/extensions.py b/lib/spack/spack/cmd/extensions.py
index b2433dfae4..5ba0281462 100644
--- a/lib/spack/spack/cmd/extensions.py
+++ b/lib/spack/spack/cmd/extensions.py
@@ -9,8 +9,7 @@ import llnl.util.tty as tty
 from llnl.util.tty.colify import colify
 
 import spack.environment as ev
-import spack.cmd
-import spack.cmd.find
+import spack.cmd as cmd
 import spack.repo
 import spack.store
 from spack.filesystem_view import YamlFilesystemView
@@ -21,16 +20,16 @@ level = "long"
 
 
 def setup_parser(subparser):
-    format_group = subparser.add_mutually_exclusive_group()
-    format_group.add_argument(
+    subparser.add_argument(
         '-l', '--long', action='store_true', dest='long',
         help='show dependency hashes as well as versions')
-    format_group.add_argument(
-        '-p', '--paths', action='store_const', dest='mode', const='paths',
-        help='show paths to extension install directories')
-    format_group.add_argument(
-        '-d', '--deps', action='store_const', dest='mode', const='deps',
-        help='show full dependency DAG of extensions')
+
+    subparser.add_argument('-d', '--deps', action='store_true',
+                           help='output dependencies along with found specs')
+
+    subparser.add_argument('-p', '--paths', action='store_true',
+                           help='show paths to package install directories')
+
     subparser.add_argument(
         '-s', '--show', dest='show', metavar='TYPE', type=str,
         default='all',
@@ -69,7 +68,7 @@ def extensions(parser, args):
     #
     # Checks
     #
-    spec = spack.cmd.parse_specs(args.spec)
+    spec = cmd.parse_specs(args.spec)
     if len(spec) > 1:
         tty.die("Can only list extensions for one package.")
 
@@ -77,14 +76,11 @@ def extensions(parser, args):
         tty.die("%s is not an extendable package." % spec[0].name)
 
     env = ev.get_env(args, 'extensions')
-    spec = spack.cmd.disambiguate_spec(spec[0], env)
+    spec = cmd.disambiguate_spec(spec[0], env)
 
     if not spec.package.extendable:
         tty.die("%s does not have extensions." % spec.short_spec)
 
-    if not args.mode:
-        args.mode = 'short'
-
     if show_packages:
         #
         # List package names of extensions
@@ -116,7 +112,7 @@ def extensions(parser, args):
             tty.msg("None installed.")
         else:
             tty.msg("%d installed:" % len(installed))
-            spack.cmd.find.display_specs(installed, mode=args.mode)
+            cmd.display_specs(installed, args)
 
     if show_activated:
         #
@@ -129,5 +125,4 @@ def extensions(parser, args):
             tty.msg("None activated.")
         else:
             tty.msg("%d currently activated:" % len(activated))
-            spack.cmd.find.display_specs(
-                activated.values(), mode=args.mode, long=args.long)
+            cmd.display_specs(activated.values(), args)
diff --git a/lib/spack/spack/cmd/find.py b/lib/spack/spack/cmd/find.py
index afe8cf9b7c..6376ef6d88 100644
--- a/lib/spack/spack/cmd/find.py
+++ b/lib/spack/spack/cmd/find.py
@@ -22,17 +22,6 @@ level = "short"
 
 def setup_parser(subparser):
     format_group = subparser.add_mutually_exclusive_group()
-    format_group.add_argument('-s', '--short',
-                              action='store_const',
-                              dest='mode',
-                              const='short',
-                              default='short',
-                              help='show only specs (default)')
-    format_group.add_argument('-p', '--paths',
-                              action='store_const',
-                              dest='mode',
-                              const='paths',
-                              help='show paths to package install directories')
     format_group.add_argument(
         "--format", action="store", default=None,
         help="output specs with the specified format string")
@@ -40,15 +29,17 @@ def setup_parser(subparser):
         "--json", action="store_true", default=False,
         help="output specs as machine-readable json records")
 
-    # TODO: separate this entirely from the "mode" option -- it's
-    # TODO: orthogonal, but changing it for all commands that use it with
-    # TODO: display_spec is tricky. Also make -pd work together properly.
+    subparser.add_argument('-d', '--deps', action='store_true',
+                           help='output dependencies along with found specs')
+
+    subparser.add_argument('-p', '--paths', action='store_true',
+                           help='show paths to package install directories')
+    subparser.add_argument(
+        '--sections', action='store_true', default=None, dest='sections',
+        help='group specs in arch/compiler sections (default on)')
     subparser.add_argument(
-        '-d', '--deps',
-        action='store_const',
-        dest='mode',
-        const='deps',
-        help='output dependencies along with found specs')
+        '--no-sections', action='store_false', default=None, dest='sections',
+        help='do not group specs by arch/compiler')
 
     arguments.add_common_arguments(
         subparser, ['long', 'very_long', 'tags'])
@@ -186,6 +177,10 @@ def find(parser, args):
     if env:
         decorator, added, roots, removed = setup_env(env)
 
+    # use sections by default except with format.
+    if args.sections is None:
+        args.sections = not args.format
+
     # Exit early if no package matches the constraint
     if not results and args.constraint:
         msg = "No package matches the query: {0}"
@@ -199,13 +194,12 @@ def find(parser, args):
         results = [x for x in results if x.name in packages_with_tags]
 
     # Display the result
-    if args.format:
-        cmd.display_formatted_specs(
-            results, args.format, deps=(args.mode == "deps"))
-    elif args.json:
-        cmd.display_specs_as_json(results, deps=(args.mode == "deps"))
+    if args.json:
+        cmd.display_specs_as_json(results, deps=args.deps)
     else:
         if env:
             display_env(env, args, decorator)
-        tty.msg("%s" % plural(len(results), 'installed package'))
-        cmd.display_specs(results, args, decorator=decorator, all_headers=True)
+        if args.sections:
+            tty.msg("%s" % plural(len(results), 'installed package'))
+        cmd.display_specs(
+            results, args, decorator=decorator, all_headers=True)
diff --git a/lib/spack/spack/test/cmd/find.py b/lib/spack/spack/test/cmd/find.py
index 7ee5487cc4..7abe63cd4f 100644
--- a/lib/spack/spack/test/cmd/find.py
+++ b/lib/spack/spack/test/cmd/find.py
@@ -35,7 +35,7 @@ def mock_display(monkeypatch, specs):
     def display(x, *args, **kwargs):
         specs.extend(x)
 
-    monkeypatch.setattr(spack.cmd.find, 'display_specs', display)
+    monkeypatch.setattr(spack.cmd, 'display_specs', display)
 
 
 def test_query_arguments():