cmd/checksum: allow adding new versions to package (#24532)

This adds super-lazy maintainer mode to `spack checksum`: Instead of
only printing the new checksums to the terminal, `-a` and
`--add-to-package` will add the new checksums to the `package.py` file
and open it in the editor afterwards for final checks.

4 files changed, 99 insertions, 9 deletions

diff --git a/lib/spack/spack/cmd/checksum.py b/lib/spack/spack/cmd/checksum.py
index 0483f4dd91..9fb7525465 100644
--- a/lib/spack/spack/cmd/checksum.py
+++ b/lib/spack/spack/cmd/checksum.py
@@ -6,6 +6,7 @@
 from __future__ import print_function
 
 import argparse
+import sys
 
 import llnl.util.tty as tty
 
@@ -16,6 +17,7 @@ import spack.spec
 import spack.stage
 import spack.util.crypto
 from spack.package_base import deprecated_version, preferred_version
+from spack.util.editor import editor
 from spack.util.naming import valid_fully_qualified_module_name
 from spack.version import VersionBase, ver
 
@@ -53,6 +55,13 @@ def setup_parser(subparser):
         default=False,
         help="checksum the preferred version only",
     )
+    subparser.add_argument(
+        "-a",
+        "--add-to-package",
+        action="store_true",
+        default=False,
+        help="add new versions to package",
+    )
     arguments.add_common_arguments(subparser, ["package"])
     subparser.add_argument(
         "versions", nargs=argparse.REMAINDER, help="versions to generate checksums for"
@@ -118,3 +127,46 @@ def checksum(parser, args):
     print()
     print(version_lines)
     print()
+
+    if args.add_to_package:
+        filename = spack.repo.path.filename_for_package_name(pkg.name)
+        # Make sure we also have a newline after the last version
+        versions = [v + "\n" for v in version_lines.splitlines()]
+        versions.append("\n")
+        # We need to insert the versions in reversed order
+        versions.reverse()
+        versions.append("    # FIXME: Added by `spack checksum`\n")
+        version_line = None
+
+        with open(filename, "r") as f:
+            lines = f.readlines()
+            for i in range(len(lines)):
+                # Black is drunk, so this is what it looks like for now
+                # See https://github.com/psf/black/issues/2156 for more information
+                if lines[i].startswith("    # FIXME: Added by `spack checksum`") or lines[
+                    i
+                ].startswith("    version("):
+                    version_line = i
+                    break
+
+        if version_line is not None:
+            for v in versions:
+                lines.insert(version_line, v)
+
+            with open(filename, "w") as f:
+                f.writelines(lines)
+
+            msg = "opening editor to verify"
+
+            if not sys.stdout.isatty():
+                msg = "please verify"
+
+            tty.info(
+                "Added {0} new versions to {1}, "
+                "{2}.".format(len(versions) - 2, args.package, msg)
+            )
+
+            if sys.stdout.isatty():
+                editor(filename)
+        else:
+            tty.warn("Could not add new versions to {0}.".format(args.package))
diff --git a/lib/spack/spack/test/cmd/checksum.py b/lib/spack/spack/test/cmd/checksum.py
index 84696bb6f8..129af2bf5f 100644
--- a/lib/spack/spack/test/cmd/checksum.py
+++ b/lib/spack/spack/test/cmd/checksum.py
@@ -20,16 +20,17 @@ spack_checksum = SpackCommand("checksum")
 @pytest.mark.parametrize(
     "arguments,expected",
     [
-        (["--batch", "patch"], (True, False, False)),
-        (["--latest", "patch"], (False, True, False)),
-        (["--preferred", "patch"], (False, False, True)),
+        (["--batch", "patch"], (True, False, False, False)),
+        (["--latest", "patch"], (False, True, False, False)),
+        (["--preferred", "patch"], (False, False, True, False)),
+        (["--add-to-package", "patch"], (False, False, False, True)),
     ],
 )
 def test_checksum_args(arguments, expected):
     parser = argparse.ArgumentParser()
     spack.cmd.checksum.setup_parser(parser)
     args = parser.parse_args(arguments)
-    check = args.batch, args.latest, args.preferred
+    check = args.batch, args.latest, args.preferred, args.add_to_package
     assert check == expected
 
 
@@ -40,9 +41,10 @@ def test_checksum_args(arguments, expected):
         (["--batch", "preferred-test"], "version of preferred-test"),
         (["--latest", "preferred-test"], "Found 1 version"),
         (["--preferred", "preferred-test"], "Found 1 version"),
+        (["--add-to-package", "preferred-test"], "Added 1 new versions to"),
     ],
 )
-def test_checksum(arguments, expected, mock_packages, mock_stage):
+def test_checksum(arguments, expected, mock_packages, mock_clone_repo, mock_stage):
     output = spack_checksum(*arguments)
     assert expected in output
     assert "version(" in output
@@ -62,19 +64,31 @@ def test_checksum_interactive(mock_packages, mock_fetch, mock_stage, monkeypatch
     assert "version(" in output
 
 
-def test_checksum_versions(mock_packages, mock_fetch, mock_stage):
+def test_checksum_versions(mock_packages, mock_clone_repo, mock_fetch, mock_stage):
     pkg_cls = spack.repo.path.get_pkg_class("preferred-test")
     versions = [str(v) for v in pkg_cls.versions if not v.isdevelop()]
     output = spack_checksum("preferred-test", versions[0])
     assert "Found 1 version" in output
     assert "version(" in output
+    output = spack_checksum("--add-to-package", "preferred-test", versions[0])
+    assert "Found 1 version" in output
+    assert "version(" in output
+    assert "Added 1 new versions to" in output
 
 
-def test_checksum_missing_version(mock_packages, mock_fetch, mock_stage):
+def test_checksum_missing_version(mock_packages, mock_clone_repo, mock_fetch, mock_stage):
     output = spack_checksum("preferred-test", "99.99.99", fail_on_error=False)
     assert "Could not find any remote versions" in output
+    output = spack_checksum("--add-to-package", "preferred-test", "99.99.99", fail_on_error=False)
+    assert "Could not find any remote versions" in output
+    assert "Added 1 new versions to" not in output
 
 
-def test_checksum_deprecated_version(mock_packages, mock_fetch, mock_stage):
+def test_checksum_deprecated_version(mock_packages, mock_clone_repo, mock_fetch, mock_stage):
     output = spack_checksum("deprecated-versions", "1.1.0", fail_on_error=False)
     assert "Version 1.1.0 is deprecated" in output
+    output = spack_checksum(
+        "--add-to-package", "deprecated-versions", "1.1.0", fail_on_error=False
+    )
+    assert "Version 1.1.0 is deprecated" in output
+    assert "Added 1 new versions to" not in output
diff --git a/lib/spack/spack/test/conftest.py b/lib/spack/spack/test/conftest.py
index c19c0c5528..0a7a21116a 100644
--- a/lib/spack/spack/test/conftest.py
+++ b/lib/spack/spack/test/conftest.py
@@ -1604,6 +1604,30 @@ repo:
     shutil.rmtree(str(repodir))
 
 
+@pytest.fixture(scope="function")
+def mock_clone_repo(tmpdir_factory):
+    """Create a cloned repository."""
+    repo_namespace = "mock_clone_repo"
+    repodir = tmpdir_factory.mktemp(repo_namespace)
+    yaml = repodir.join("repo.yaml")
+    yaml.write(
+        """
+repo:
+    namespace: mock_clone_repo
+"""
+    )
+
+    shutil.copytree(
+        os.path.join(spack.paths.mock_packages_path, spack.repo.packages_dir_name),
+        os.path.join(str(repodir), spack.repo.packages_dir_name),
+    )
+
+    with spack.repo.use_repositories(str(repodir)) as repo:
+        yield repo, repodir
+
+    shutil.rmtree(str(repodir))
+
+
 ##########
 # Class and fixture to work around problems raising exceptions in directives,
 # which cause tests like test_from_list_url to hang for Python 2.x metaclass
diff --git a/share/spack/spack-completion.bash b/share/spack/spack-completion.bash
index 24cc2e67a0..8abf64a01e 100755
--- a/share/spack/spack-completion.bash
+++ b/share/spack/spack-completion.bash
@@ -580,7 +580,7 @@ _spack_change() {
 _spack_checksum() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help --keep-stage -b --batch -l --latest -p --preferred"
+        SPACK_COMPREPLY="-h --help --keep-stage -b --batch -l --latest -p --preferred -a --add-to-package"
     else
         _all_packages
     fi