Add spack checksum --verify, fix --add (#38458)

* Add rewrite of spack checksum to include --verify and better add versions to package.py files
* Fix formatting and remove unused import
* Update checksum unit-tests to correctly test multiple versions and add to package
* Remove references to latest in stage.py
* Update bash-completion scripts to fix unit tests failures
* Fix docs generation
* Remove unused url_dict argument from methods
* Reduce chance of redundant remote_versions work
* Add print() before tty.die() to increase error readablity
* Update version regular expression to allow for multi-line versions
* Add a few unit tests to improve test coverage
* Update command completion
* Add type hints to added functions and fix a few py-lint suggestions
* Add @no_type_check to prevent mypy from failing on pkg.versions
* Add type hints to format.py and fix unit test
* Black format lib/spack/spack/package_base.py
* Attempt ignoring type errors
* Add optional dict type hint and declare versions in PackageBase
* Refactor util/format.py to allow for url_dict as an optional parameter
* Directly reference PackageBase class instead of using TypeVar
* Fix comment typo

---------

Co-authored-by: Tamara Dahlgren <dahlgren1@llnl.gov>

lib/spack/spack/cmd/checksum.py

@@ -4,18 +4,21 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
 import argparse
+import re
 import sys
 
-import llnl.util.tty as tty
+import llnl.util.lang
+from llnl.util import tty
 
 import spack.cmd
-import spack.cmd.common.arguments as arguments
 import spack.repo
 import spack.spec
 import spack.stage
 import spack.util.crypto
-from spack.package_base import deprecated_version, preferred_version
+from spack.cmd.common import arguments
+from spack.package_base import PackageBase, deprecated_version, preferred_version
 from spack.util.editor import editor
+from spack.util.format import get_version_lines
 from spack.util.naming import valid_fully_qualified_module_name
 from spack.version import Version
 
@@ -31,35 +34,38 @@ def setup_parser(subparser):
         default=False,
         help="don't clean up staging area when command completes",
     )
-    sp = subparser.add_mutually_exclusive_group()
+    subparser.add_argument(
-    sp.add_argument(
         "-b",
         "--batch",
         action="store_true",
         default=False,
         help="don't ask which versions to checksum",
     )
-    sp.add_argument(
+    subparser.add_argument(
         "-l",
         "--latest",
         action="store_true",
         default=False,
-        help="checksum the latest available version only",
+        help="checksum the latest available version",
     )
-    sp.add_argument(
+    subparser.add_argument(
         "-p",
         "--preferred",
         action="store_true",
         default=False,
-        help="checksum the preferred version only",
+        help="checksum the known Spack preferred version",
     )
-    subparser.add_argument(
+    modes_parser = subparser.add_mutually_exclusive_group()
+    modes_parser.add_argument(
         "-a",
         "--add-to-package",
         action="store_true",
         default=False,
         help="add new versions to package",
     )
+    modes_parser.add_argument(
+        "--verify", action="store_true", default=False, help="verify known package checksums"
+    )
     arguments.add_common_arguments(subparser, ["package"])
     subparser.add_argument(
         "versions", nargs=argparse.REMAINDER, help="versions to generate checksums for"
@@ -80,86 +86,171 @@ def checksum(parser, args):
     pkg_cls = spack.repo.path.get_pkg_class(args.package)
     pkg = pkg_cls(spack.spec.Spec(args.package))
 
+    # Build a list of versions to checksum
+    versions = [Version(v) for v in args.versions]
+
+    # Define placeholder for remote versions.
+    # This'll help reduce redundant work if we need to check for the existance
+    # of remote versions more than once.
+    remote_versions = None
+
+    # Add latest version if requested
+    if args.latest:
+        remote_versions = pkg.fetch_remote_versions()
+        if len(remote_versions) > 0:
+            latest_version = sorted(remote_versions.keys(), reverse=True)[0]
+            versions.append(latest_version)
+
+    # Add preferred version if requested
+    if args.preferred:
+        versions.append(preferred_version(pkg))
+
+    # Store a dict of the form version -> URL
     url_dict = {}
-    if not args.versions and args.preferred:
-        versions = [preferred_version(pkg)]
-    else:
-        versions = [Version(v) for v in args.versions]
 
-    if versions:
+    for version in versions:
-        remote_versions = None
+        if deprecated_version(pkg, version):
-        for version in versions:
+            tty.warn(f"Version {version} is deprecated")
-            if deprecated_version(pkg, version):
-                tty.warn("Version {0} is deprecated".format(version))
 
-            url = pkg.find_valid_url_for_version(version)
+        url = pkg.find_valid_url_for_version(version)
-            if url is not None:
+        if url is not None:
-                url_dict[version] = url
+            url_dict[version] = url
-                continue
+            continue
-            # if we get here, it's because no valid url was provided by the package
+        # if we get here, it's because no valid url was provided by the package
-            # do expensive fallback to try to recover
+        # do expensive fallback to try to recover
-            if remote_versions is None:
+        if remote_versions is None:
-                remote_versions = pkg.fetch_remote_versions()
+            remote_versions = pkg.fetch_remote_versions()
-            if version in remote_versions:
+        if version in remote_versions:
-                url_dict[version] = remote_versions[version]
+            url_dict[version] = remote_versions[version]
-    else:
+
-        url_dict = pkg.fetch_remote_versions()
+    if len(versions) <= 0:
+        if remote_versions is None:
+            remote_versions = pkg.fetch_remote_versions()
+        url_dict = remote_versions
 
     if not url_dict:
-        tty.die("Could not find any remote versions for {0}".format(pkg.name))
+        tty.die(f"Could not find any remote versions for {pkg.name}")
 
-    version_lines = spack.stage.get_checksums_for_versions(
+    # print an empty line to create a new output section block
+    print()
+
+    version_hashes = spack.stage.get_checksums_for_versions(
         url_dict,
         pkg.name,
         keep_stage=args.keep_stage,
-        batch=(args.batch or len(args.versions) > 0 or len(url_dict) == 1),
+        batch=(args.batch or len(versions) > 0 or len(url_dict) == 1),
-        latest=args.latest,
         fetch_options=pkg.fetch_options,
     )
 
+    if args.verify:
+        print_checksum_status(pkg, version_hashes)
+        sys.exit(0)
+
+    # convert dict into package.py version statements
+    version_lines = get_version_lines(version_hashes, url_dict)
     print()
     print(version_lines)
     print()
 
     if args.add_to_package:
-        filename = spack.repo.path.filename_for_package_name(pkg.name)
+        add_versions_to_package(pkg, version_lines)
-        # 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:
+def print_checksum_status(pkg: PackageBase, version_hashes: dict):
-            for v in versions:
+    """
-                lines.insert(version_line, v)
+    Verify checksums present in version_hashes against those present
+    in the package's instructions.
 
-            with open(filename, "w") as f:
+    Args:
-                f.writelines(lines)
+        pkg (spack.package_base.PackageBase): A package class for a given package in Spack.
+        version_hashes (dict): A dictionary of the form: version -> checksum.
 
-            msg = "opening editor to verify"
+    """
+    results = []
+    num_verified = 0
+    failed = False
 
-            if not sys.stdout.isatty():
+    max_len = max(len(str(v)) for v in version_hashes)
-                msg = "please verify"
+    num_total = len(version_hashes)
 
-            tty.info(
+    for version, sha in version_hashes.items():
-                "Added {0} new versions to {1}, "
+        if version not in pkg.versions:
-                "{2}.".format(len(versions) - 2, args.package, msg)
+            msg = "No previous checksum"
-            )
+            status = "-"
+
+        elif sha == pkg.versions[version]["sha256"]:
+            msg = "Correct"
+            status = "="
+            num_verified += 1
 
-            if sys.stdout.isatty():
-                editor(filename)
         else:
-            tty.warn("Could not add new versions to {0}.".format(args.package))
+            msg = sha
+            status = "x"
+            failed = True
+
+        results.append("{0:{1}}  {2} {3}".format(str(version), max_len, f"[{status}]", msg))
+
+    # Display table of checksum results.
+    tty.msg(f"Verified {num_verified} of {num_total}", "", *llnl.util.lang.elide_list(results), "")
+
+    # Terminate at the end of function to prevent additional output.
+    if failed:
+        print()
+        tty.die("Invalid checksums found.")
+
+
+def add_versions_to_package(pkg: PackageBase, version_lines: str):
+    """
+    Add checksumed versions to a package's instructions and open a user's
+    editor so they may double check the work of the function.
+
+    Args:
+        pkg (spack.package_base.PackageBase): A package class for a given package in Spack.
+        version_lines (str): A string of rendered version lines.
+
+    """
+    # Get filename and path for package
+    filename = spack.repo.path.filename_for_package_name(pkg.name)
+    num_versions_added = 0
+
+    version_statement_re = re.compile(r"([\t ]+version\([^\)]*\))")
+    version_re = re.compile(r'[\t ]+version\(\s*"([^"]+)"[^\)]*\)')
+
+    # Split rendered version lines into tuple of (version, version_line)
+    # We reverse sort here to make sure the versions match the version_lines
+    new_versions = []
+    for ver_line in version_lines.split("\n"):
+        match = version_re.match(ver_line)
+        if match:
+            new_versions.append((Version(match.group(1)), ver_line))
+
+    with open(filename, "r+") as f:
+        contents = f.read()
+        split_contents = version_statement_re.split(contents)
+
+        for i, subsection in enumerate(split_contents):
+            # If there are no more versions to add we should exit
+            if len(new_versions) <= 0:
+                break
+
+            # Check if the section contains a version
+            contents_version = version_re.match(subsection)
+            if contents_version is not None:
+                parsed_version = Version(contents_version.group(1))
+
+                if parsed_version < new_versions[0][0]:
+                    split_contents[i:i] = [new_versions.pop(0)[1], " # FIX ME", "\n"]
+                    num_versions_added += 1
+
+                elif parsed_version == new_versions[0][0]:
+                    new_versions.pop(0)
+
+        # Seek back to the start of the file so we can rewrite the file contents.
+        f.seek(0)
+        f.writelines("".join(split_contents))
+
+        tty.msg(f"Added {num_versions_added} new versions to {pkg.name}")
+        tty.msg(f"Open {filename} to review the additions.")
+
+    if sys.stdout.isatty():
+        editor(filename)

lib/spack/spack/cmd/create.py

@@ -17,6 +17,7 @@
 from spack.url import UndetectableNameError, UndetectableVersionError, parse_name, parse_version
 from spack.util.editor import editor
 from spack.util.executable import ProcessError, which
+from spack.util.format import get_version_lines
 from spack.util.naming import mod_to_class, simplify_name, valid_fully_qualified_module_name
 
 description = "create a new package file"
@@ -832,13 +833,15 @@ def get_versions(args, name):
             version = parse_version(args.url)
             url_dict = {version: args.url}
 
-        versions = spack.stage.get_checksums_for_versions(
+        version_hashes = spack.stage.get_checksums_for_versions(
             url_dict,
             name,
             first_stage_function=guesser,
             keep_stage=args.keep_stage,
             batch=(args.batch or len(url_dict) == 1),
         )
+
+        versions = get_version_lines(version_hashes, url_dict)
     else:
         versions = unhashed_versions
 

lib/spack/spack/package_base.py

@@ -514,6 +514,10 @@ class PackageBase(WindowsRPath, PackageViewMixin, metaclass=PackageMeta):
     # These are default values for instance variables.
     #
 
+    # Declare versions dictionary as placeholder for values.
+    # This allows analysis tools to correctly interpret the class attributes.
+    versions: dict
+
     #: By default, packages are not virtual
     #: Virtual packages override this attribute
     virtual = False

lib/spack/spack/stage.py

@@ -884,23 +884,19 @@ def get_checksums_for_versions(url_dict, name, **kwargs):
         keep_stage (bool): whether to keep staging area when command completes
         batch (bool): whether to ask user how many versions to fetch (false)
             or fetch all versions (true)
-        latest (bool): whether to take the latest version (true) or all (false)
         fetch_options (dict): Options used for the fetcher (such as timeout
             or cookies)
 
     Returns:
-        (str): A multi-line string containing versions and corresponding hashes
+        (dict): A dictionary of the form: version -> checksum
 
     """
     batch = kwargs.get("batch", False)
     fetch_options = kwargs.get("fetch_options", None)
     first_stage_function = kwargs.get("first_stage_function", None)
     keep_stage = kwargs.get("keep_stage", False)
-    latest = kwargs.get("latest", False)
 
     sorted_versions = sorted(url_dict.keys(), reverse=True)
-    if latest:
-        sorted_versions = sorted_versions[:1]
 
     # Find length of longest string in the list for padding
     max_len = max(len(str(v)) for v in sorted_versions)
@@ -915,7 +911,7 @@ def get_checksums_for_versions(url_dict, name, **kwargs):
     )
     print()
 
-    if batch or latest:
+    if batch:
         archives_to_fetch = len(sorted_versions)
     else:
         archives_to_fetch = tty.get_number(
@@ -929,14 +925,10 @@ def get_checksums_for_versions(url_dict, name, **kwargs):
     urls = [url_dict[v] for v in versions]
 
     tty.debug("Downloading...")
-    version_hashes = []
+    version_hashes = {}
     i = 0
     errors = []
     for url, version in zip(urls, versions):
-        # Wheels should not be expanded during staging
-        expand_arg = ""
-        if url.endswith(".whl") or ".whl#" in url:
-            expand_arg = ", expand=False"
         try:
             if fetch_options:
                 url_or_fs = fs.URLFetchStrategy(url, fetch_options=fetch_options)
@@ -951,8 +943,8 @@ def get_checksums_for_versions(url_dict, name, **kwargs):
                     first_stage_function(stage, url)
 
                 # Checksum the archive and add it to the list
-                version_hashes.append(
+                version_hashes[version] = spack.util.crypto.checksum(
-                    (version, spack.util.crypto.checksum(hashlib.sha256, stage.archive_file))
+                    hashlib.sha256, stage.archive_file
                 )
                 i += 1
         except FailedDownloadError:
@@ -966,17 +958,12 @@ def get_checksums_for_versions(url_dict, name, **kwargs):
     if not version_hashes:
         tty.die("Could not fetch any versions for {0}".format(name))
 
-    # Generate the version directives to put in a package.py
-    version_lines = "\n".join(
-        ['    version("{0}", sha256="{1}"{2})'.format(v, h, expand_arg) for v, h in version_hashes]
-    )
-
     num_hash = len(version_hashes)
     tty.debug(
         "Checksummed {0} version{1} of {2}:".format(num_hash, "" if num_hash == 1 else "s", name)
     )
 
-    return version_lines
+    return version_hashes
 
 
 class StageError(spack.error.SpackError):

lib/spack/spack/test/cmd/checksum.py

@@ -12,6 +12,7 @@
 
 import spack.cmd.checksum
 import spack.repo
+import spack.spec
 from spack.main import SpackCommand
 
 spack_checksum = SpackCommand("checksum")
@@ -20,17 +21,18 @@
 @pytest.mark.parametrize(
     "arguments,expected",
     [
-        (["--batch", "patch"], (True, False, False, False)),
+        (["--batch", "patch"], (True, False, False, False, False)),
-        (["--latest", "patch"], (False, True, False, False)),
+        (["--latest", "patch"], (False, True, False, False, False)),
-        (["--preferred", "patch"], (False, False, True, False)),
+        (["--preferred", "patch"], (False, False, True, False, False)),
-        (["--add-to-package", "patch"], (False, False, False, True)),
+        (["--add-to-package", "patch"], (False, False, False, True, False)),
+        (["--verify", "patch"], (False, 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, args.add_to_package
+    check = args.batch, args.latest, args.preferred, args.add_to_package, args.verify
     assert check == expected
 
 
@@ -41,13 +43,18 @@ 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"),
+        (["--add-to-package", "preferred-test"], "Added 0 new versions to"),
+        (["--verify", "preferred-test"], "Verified 1 of 1"),
+        (["--verify", "zlib", "1.2.13"], "1.2.13  [-] No previous checksum"),
     ],
 )
 def test_checksum(arguments, expected, mock_packages, mock_clone_repo, mock_stage):
     output = spack_checksum(*arguments)
     assert expected in output
-    assert "version(" in output
+
+    # --verify doesn't print versions strings like other flags
+    if "--verify" not in arguments:
+        assert "version(" in output
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="Not supported on Windows (yet)")
@@ -65,15 +72,14 @@ def _get_number(*args, **kwargs):
 
 
 def test_checksum_versions(mock_packages, mock_clone_repo, mock_fetch, mock_stage):
-    pkg_cls = spack.repo.path.get_pkg_class("preferred-test")
+    pkg_cls = spack.repo.path.get_pkg_class("zlib")
-    versions = [str(v) for v in pkg_cls.versions if not v.isdevelop()]
+    versions = [str(v) for v in pkg_cls.versions]
-    output = spack_checksum("preferred-test", versions[0])
+    output = spack_checksum("zlib", *versions)
-    assert "Found 1 version" in output
+    assert "Found 3 versions" in output
     assert "version(" in output
-    output = spack_checksum("--add-to-package", "preferred-test", versions[0])
+    output = spack_checksum("--add-to-package", "zlib", *versions)
-    assert "Found 1 version" in output
+    assert "Found 3 versions" in output
-    assert "version(" in output
+    assert "Added 0 new versions to" in output
-    assert "Added 1 new versions to" in output
 
 
 def test_checksum_missing_version(mock_packages, mock_clone_repo, mock_fetch, mock_stage):
@@ -91,4 +97,30 @@ def test_checksum_deprecated_version(mock_packages, mock_clone_repo, mock_fetch,
         "--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
+    assert "Added 0 new versions to" not in output
+
+
+def test_checksum_at(mock_packages):
+    pkg_cls = spack.repo.path.get_pkg_class("zlib")
+    versions = [str(v) for v in pkg_cls.versions]
+    output = spack_checksum(f"zlib@{versions[0]}")
+    assert "Found 1 version" in output
+
+
+def test_checksum_url(mock_packages):
+    pkg_cls = spack.repo.path.get_pkg_class("zlib")
+    output = spack_checksum(f"{pkg_cls.url}", fail_on_error=False)
+    assert "accepts package names" in output
+
+
+def test_checksum_verification_fails(install_mockery, capsys):
+    spec = spack.spec.Spec("zlib").concretized()
+    pkg = spec.package
+    versions = list(pkg.versions.keys())
+    version_hashes = {versions[0]: "abadhash", spack.version.Version("0.1"): "123456789"}
+    with pytest.raises(SystemExit):
+        spack.cmd.checksum.print_checksum_status(pkg, version_hashes)
+    out = str(capsys.readouterr())
+    assert out.count("Correct") == 0
+    assert "No previous checksum" in out
+    assert "Invalid checksum" in out

lib/spack/spack/util/format.py

@@ -0,0 +1,36 @@
+# Copyright 2013-2023 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+from typing import Optional
+
+
+def get_version_lines(version_hashes_dict: dict, url_dict: Optional[dict] = None) -> str:
+    """
+    Renders out a set of versions like those found in a package's
+    package.py file for a given set of versions and hashes.
+
+    Args:
+        version_hashes_dict (dict): A dictionary of the form: version -> checksum.
+        url_dict (dict): A dictionary of the form: version -> URL.
+
+    Returns:
+        (str): Rendered version lines.
+    """
+    version_lines = []
+
+    for v, h in version_hashes_dict.items():
+        expand_arg = ""
+
+        # Extract the url for a version if url_dict is provided.
+        url = ""
+        if url_dict is not None and v in url_dict:
+            url = url_dict[v]
+
+        # Add expand_arg since wheels should not be expanded during stanging
+        if url.endswith(".whl") or ".whl#" in url:
+            expand_arg = ", expand=False"
+        version_lines.append(f'    version("{v}", sha256="{h}"{expand_arg})')
+
+    return "\n".join(version_lines)

share/spack/spack-completion.bash

@@ -608,7 +608,7 @@ _spack_change() {
 _spack_checksum() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help --keep-stage -b --batch -l --latest -p --preferred -a --add-to-package"
+        SPACK_COMPREPLY="-h --help --keep-stage -b --batch -l --latest -p --preferred -a --add-to-package --verify"
     else
         _all_packages
     fi

share/spack/spack-completion.fish

@@ -882,7 +882,7 @@ complete -c spack -n '__fish_spack_using_command change' -s a -l all -f -a all
 complete -c spack -n '__fish_spack_using_command change' -s a -l all -d 'change all matching specs (allow changing more than one spec)'
 
 # spack checksum
-set -g __fish_spack_optspecs_spack_checksum h/help keep-stage b/batch l/latest p/preferred a/add-to-package
+set -g __fish_spack_optspecs_spack_checksum h/help keep-stage b/batch l/latest p/preferred a/add-to-package verify
 complete -c spack -n '__fish_spack_using_command_pos 0 checksum' -f -a '(__fish_spack_packages)'
 complete -c spack -n '__fish_spack_using_command_pos_remainder 1 checksum' -f -a '(__fish_spack_package_versions $__fish_spack_argparse_argv[1])'
 complete -c spack -n '__fish_spack_using_command checksum' -s h -l help -f -a help
@@ -892,11 +892,13 @@ complete -c spack -n '__fish_spack_using_command checksum' -l keep-stage -d 'don
 complete -c spack -n '__fish_spack_using_command checksum' -s b -l batch -f -a batch
 complete -c spack -n '__fish_spack_using_command checksum' -s b -l batch -d 'don\'t ask which versions to checksum'
 complete -c spack -n '__fish_spack_using_command checksum' -s l -l latest -f -a latest
-complete -c spack -n '__fish_spack_using_command checksum' -s l -l latest -d 'checksum the latest available version only'
+complete -c spack -n '__fish_spack_using_command checksum' -s l -l latest -d 'checksum the latest available version'
 complete -c spack -n '__fish_spack_using_command checksum' -s p -l preferred -f -a preferred
-complete -c spack -n '__fish_spack_using_command checksum' -s p -l preferred -d 'checksum the preferred version only'
+complete -c spack -n '__fish_spack_using_command checksum' -s p -l preferred -d 'checksum the known Spack preferred version'
 complete -c spack -n '__fish_spack_using_command checksum' -s a -l add-to-package -f -a add_to_package
 complete -c spack -n '__fish_spack_using_command checksum' -s a -l add-to-package -d 'add new versions to package'
+complete -c spack -n '__fish_spack_using_command checksum' -l verify -f -a verify
+complete -c spack -n '__fish_spack_using_command checksum' -l verify -d 'verify known package checksums'
 
 # spack ci
 set -g __fish_spack_optspecs_spack_ci h/help