new command: `spack pkg grep` to search package files (#34388)

It's very common for us to tell users to grep through the existing Spack packages to
find examples of what they want, and it's also very common for package developers to do
it. Now, searching packages is even easier.

`spack pkg grep` runs grep on all `package.py` files in repos known to Spack. It has no
special options other than the search string; all options passed to it are forwarded
along to `grep`.

```console
> spack pkg grep --help
usage: spack pkg grep [--help] ...

positional arguments:
  grep_args  arguments for grep

options:
  --help     show this help message and exit
```

```console
> spack pkg grep CMakePackage | head -3
/Users/gamblin2/src/spack/var/spack/repos/builtin/packages/3dtk/package.py:class _3dtk(CMakePackage):
/Users/gamblin2/src/spack/var/spack/repos/builtin/packages/abseil-cpp/package.py:class AbseilCpp(CMakePackage):
/Users/gamblin2/src/spack/var/spack/repos/builtin/packages/accfft/package.py:class Accfft(CMakePackage, CudaPackage):
```

```console
> spack pkg grep -Eho '(\S*)\(PythonPackage\)' | head -3
AwsParallelcluster(PythonPackage)
Awscli(PythonPackage)
Bueno(PythonPackage)
```

## Return Value

This retains the return value semantics of `grep`:
* 0  for found,
* 1  for not found
* >1 for error

## Choosing a `grep`

You can set the ``SPACK_GREP`` environment variable to choose the ``grep``
executable this command should use.

3 files changed, 125 insertions, 5 deletions

diff --git a/lib/spack/spack/cmd/pkg.py b/lib/spack/spack/cmd/pkg.py
index 6de7a4bcc1..8302e3e0c1 100644
--- a/lib/spack/spack/cmd/pkg.py
+++ b/lib/spack/spack/cmd/pkg.py
@@ -5,6 +5,9 @@
 
 from __future__ import print_function
 
+import argparse
+import itertools
+import os
 import sys
 
 import llnl.util.tty as tty
@@ -14,6 +17,7 @@ import spack.cmd
 import spack.cmd.common.arguments as arguments
 import spack.paths
 import spack.repo
+import spack.util.executable as exe
 import spack.util.package_hash as ph
 
 description = "query packages associated with particular git revisions"
@@ -65,6 +69,14 @@ def setup_parser(subparser):
         "rev2", nargs="?", default="HEAD", help="revision to compare to rev1 (default is HEAD)"
     )
 
+    # explicitly add help for `spack pkg grep` with just `--help` and NOT `-h`. This is so
+    # that the very commonly used -h (no filename) argument can be passed through to grep
+    grep_parser = sp.add_parser("grep", help=pkg_grep.__doc__, add_help=False)
+    grep_parser.add_argument(
+        "grep_args", nargs=argparse.REMAINDER, default=None, help="arguments for grep"
+    )
+    grep_parser.add_argument("--help", action="help", help="show this help message and exit")
+
     source_parser = sp.add_parser("source", help=pkg_source.__doc__)
     source_parser.add_argument(
         "-c",
@@ -157,18 +169,88 @@ def pkg_hash(args):
         print(ph.package_hash(spec))
 
 
-def pkg(parser, args):
+def get_grep(required=False):
+    """Get a grep command to use with ``spack pkg grep``."""
+    return exe.which(os.environ.get("SPACK_GREP") or "grep", required=required)
+
+
+def pkg_grep(args, unknown_args):
+    """grep for strings in package.py files from all repositories"""
+    grep = get_grep(required=True)
+
+    # add a little color to the output if we can
+    if "GNU" in grep("--version", output=str):
+        grep.add_default_arg("--color=auto")
+
+    # determines number of files to grep at a time
+    grouper = lambda e: e[0] // 500
+
+    # set up iterator and save the first group to ensure we don't end up with a group of size 1
+    groups = itertools.groupby(enumerate(spack.repo.path.all_package_paths()), grouper)
+    if not groups:
+        return 0  # no packages to search
+
+    # You can force GNU grep to show filenames on every line with -H, but not POSIX grep.
+    # POSIX grep only shows filenames when you're grepping 2 or more files.  Since we
+    # don't know which one we're running, we ensure there are always >= 2 files by
+    # saving the prior group of paths and adding it to a straggling group of 1 if needed.
+    # This works unless somehow there is only one package in all of Spack.
+    _, first_group = next(groups)
+    prior_paths = [path for _, path in first_group]
+
+    # grep returns 1 for nothing found, 0 for something found, and > 1 for error
+    return_code = 1
+
+    # assemble args and run grep on a group of paths
+    def grep_group(paths):
+        all_args = args.grep_args + unknown_args + paths
+        grep(*all_args, fail_on_error=False)
+        return grep.returncode
+
+    for _, group in groups:
+        paths = [path for _, path in group]  # extract current path group
+
+        if len(paths) == 1:
+            # Only the very last group can have length 1. If it does, combine
+            # it with the prior group to ensure more than one path is grepped.
+            prior_paths += paths
+        else:
+            # otherwise run grep on the prior group
+            error = grep_group(prior_paths)
+            if error != 1:
+                return_code = error
+                if error > 1:  # fail fast on error
+                    return error
+
+            prior_paths = paths
+
+    # Handle the last remaining group after the loop
+    error = grep_group(prior_paths)
+    if error != 1:
+        return_code = error
+
+    return return_code
+
+
+def pkg(parser, args, unknown_args):
     if not spack.cmd.spack_is_git_repo():
         tty.die("This spack is not a git clone. Can't use 'spack pkg'")
 
     action = {
         "add": pkg_add,
+        "added": pkg_added,
+        "changed": pkg_changed,
         "diff": pkg_diff,
+        "hash": pkg_hash,
         "list": pkg_list,
         "removed": pkg_removed,
-        "added": pkg_added,
-        "changed": pkg_changed,
         "source": pkg_source,
-        "hash": pkg_hash,
     }
-    action[args.pkg_command](args)
+
+    # grep is special as it passes unknown arguments through
+    if args.pkg_command == "grep":
+        return pkg_grep(args, unknown_args)
+    elif unknown_args:
+        tty.die("unrecognized arguments: %s" % " ".join(unknown_args))
+    else:
+        return action[args.pkg_command](args)
diff --git a/lib/spack/spack/repo.py b/lib/spack/spack/repo.py
index 6538ab0162..2710b04920 100644
--- a/lib/spack/spack/repo.py
+++ b/lib/spack/spack/repo.py
@@ -754,6 +754,14 @@ class RepoPath(object):
     def all_package_names(self, include_virtuals=False):
         return self._all_package_names(include_virtuals)
 
+    def package_path(self, name):
+        """Get path to package.py file for this repo."""
+        return self.repo_for_pkg(name).package_path(name)
+
+    def all_package_paths(self):
+        for name in self.all_package_names():
+            yield self.package_path(name)
+
     def packages_with_tags(self, *tags):
         r = set()
         for repo in self.repos:
@@ -1153,6 +1161,14 @@ class Repo(object):
             return names
         return [x for x in names if not self.is_virtual(x)]
 
+    def package_path(self, name):
+        """Get path to package.py file for this repo."""
+        return os.path.join(self.root, packages_dir_name, name, package_file_name)
+
+    def all_package_paths(self):
+        for name in self.all_package_names():
+            yield self.package_path(name)
+
     def packages_with_tags(self, *tags):
         v = set(self.all_package_names())
         index = self.tag_index
diff --git a/lib/spack/spack/test/cmd/pkg.py b/lib/spack/spack/test/cmd/pkg.py
index 9c2b700816..2f1a1a6f3a 100644
--- a/lib/spack/spack/test/cmd/pkg.py
+++ b/lib/spack/spack/test/cmd/pkg.py
@@ -13,6 +13,7 @@ import pytest
 
 from llnl.util.filesystem import mkdirp, working_dir
 
+import spack.cmd.pkg
 import spack.main
 import spack.repo
 from spack.util.executable import which
@@ -293,3 +294,24 @@ def test_pkg_hash(mock_packages):
 
     output = pkg("hash", "multimethod").strip().split()
     assert len(output) == 1 and all(len(elt) == 32 for elt in output)
+
+
+@pytest.mark.skipif(not spack.cmd.pkg.get_grep(), reason="grep is not installed")
+def test_pkg_grep(mock_packages, capsys):
+    # only splice-* mock packages have the string "splice" in them
+    with capsys.disabled():
+        output = pkg("grep", "-l", "splice", output=str)
+
+    assert output.strip() == "\n".join(
+        spack.repo.path.get_pkg_class(name).module.__file__
+        for name in ["splice-a", "splice-h", "splice-t", "splice-vh", "splice-z"]
+    )
+
+    # ensure that this string isn't fouhnd
+    output = pkg("grep", "abcdefghijklmnopqrstuvwxyz", output=str, fail_on_error=False)
+    assert pkg.returncode == 1
+    assert output.strip() == ""
+
+    # ensure that we return > 1 for an error
+    pkg("grep", "--foobarbaz-not-an-option", output=str, fail_on_error=False)
+    assert pkg.returncode == 2