Add --preferred and --latest to`spack checksum` (#25830)

6 files changed, 104 insertions, 15 deletions

diff --git a/lib/spack/spack/cmd/checksum.py b/lib/spack/spack/cmd/checksum.py
index 96e36b39b6..b2a2bc5891 100644
--- a/lib/spack/spack/cmd/checksum.py
+++ b/lib/spack/spack/cmd/checksum.py
@@ -14,6 +14,7 @@ import spack.cmd.common.arguments as arguments
 import spack.repo
 import spack.stage
 import spack.util.crypto
+from spack.package import preferred_version
 from spack.util.naming import valid_fully_qualified_module_name
 from spack.version import Version, ver
 
@@ -26,9 +27,16 @@ def setup_parser(subparser):
     subparser.add_argument(
         '--keep-stage', action='store_true',
         help="don't clean up staging area when command completes")
-    subparser.add_argument(
+    sp = subparser.add_mutually_exclusive_group()
+    sp.add_argument(
         '-b', '--batch', action='store_true',
         help="don't ask which versions to checksum")
+    sp.add_argument(
+        '-l', '--latest', action='store_true',
+        help="checksum the latest available version only")
+    sp.add_argument(
+        '-p', '--preferred', action='store_true',
+        help="checksum the preferred version only")
     arguments.add_common_arguments(subparser, ['package'])
     subparser.add_argument(
         'versions', nargs=argparse.REMAINDER,
@@ -48,15 +56,18 @@ def checksum(parser, args):
     # Get the package we're going to generate checksums for
     pkg = spack.repo.get(args.package)
 
+    url_dict = {}
     if args.versions:
         # If the user asked for specific versions, use those
-        url_dict = {}
         for version in args.versions:
             version = ver(version)
             if not isinstance(version, Version):
                 tty.die("Cannot generate checksums for version lists or "
                         "version ranges. Use unambiguous versions.")
             url_dict[version] = pkg.url_for_version(version)
+    elif args.preferred:
+        version = preferred_version(pkg)
+        url_dict = dict([(version, pkg.url_for_version(version))])
     else:
         # Otherwise, see what versions we can find online
         url_dict = pkg.fetch_remote_versions()
@@ -76,7 +87,7 @@ def checksum(parser, args):
     version_lines = 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),
-        fetch_options=pkg.fetch_options)
+        latest=args.latest, fetch_options=pkg.fetch_options)
 
     print()
     print(version_lines)
diff --git a/lib/spack/spack/cmd/info.py b/lib/spack/spack/cmd/info.py
index 7119a23dbc..b4994a7b2d 100644
--- a/lib/spack/spack/cmd/info.py
+++ b/lib/spack/spack/cmd/info.py
@@ -17,6 +17,7 @@ import spack.cmd.common.arguments as arguments
 import spack.fetch_strategy as fs
 import spack.repo
 import spack.spec
+from spack.package import preferred_version
 
 description = 'get detailed information on a particular package'
 section = 'basic'
@@ -197,13 +198,7 @@ def print_text_info(pkg):
     else:
         pad = padder(pkg.versions, 4)
 
-        # Here we sort first on the fact that a version is marked
-        # as preferred in the package, then on the fact that the
-        # version is not develop, then lexicographically
-        key_fn = lambda v: (pkg.versions[v].get('preferred', False),
-                            not v.isdevelop(),
-                            v)
-        preferred = sorted(pkg.versions, key=key_fn).pop()
+        preferred = preferred_version(pkg)
         url = ''
         if pkg.has_code:
             url = fs.for_package_version(pkg, preferred)
diff --git a/lib/spack/spack/package.py b/lib/spack/spack/package.py
index 4fbf24c56e..5f66d8255d 100644
--- a/lib/spack/spack/package.py
+++ b/lib/spack/spack/package.py
@@ -83,6 +83,22 @@ _spack_times_log = 'install_times.json'
 _spack_configure_argsfile = 'spack-configure-args.txt'
 
 
+def preferred_version(pkg):
+    """
+    Returns a sorted list of the preferred versions of the package.
+
+    Arguments:
+        pkg (Package): The package whose versions are to be assessed.
+    """
+    # Here we sort first on the fact that a version is marked
+    # as preferred in the package, then on the fact that the
+    # version is not develop, then lexicographically
+    key_fn = lambda v: (pkg.versions[v].get('preferred', False),
+                        not v.isdevelop(),
+                        v)
+    return sorted(pkg.versions, key=key_fn).pop()
+
+
 class InstallPhase(object):
     """Manages a single phase of the installation.
 
diff --git a/lib/spack/spack/stage.py b/lib/spack/spack/stage.py
index 28feb05365..bffd06ab73 100644
--- a/lib/spack/spack/stage.py
+++ b/lib/spack/spack/stage.py
@@ -824,9 +824,7 @@ def purge():
                 remove_linked_tree(stage_path)
 
 
-def get_checksums_for_versions(
-        url_dict, name, first_stage_function=None, keep_stage=False,
-        fetch_options=None, batch=False):
+def get_checksums_for_versions(url_dict, name, **kwargs):
     """Fetches and checksums archives from URLs.
 
     This function is called by both ``spack checksum`` and ``spack
@@ -842,6 +840,7 @@ def get_checksums_for_versions(
         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)
 
@@ -849,7 +848,15 @@ def get_checksums_for_versions(
         (str): A multi-line string containing versions and corresponding hashes
 
     """
+    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)
@@ -863,7 +870,7 @@ def get_checksums_for_versions(
                  for v in sorted_versions]))
     print()
 
-    if batch:
+    if batch or latest:
         archives_to_fetch = len(sorted_versions)
     else:
         archives_to_fetch = tty.get_number(
diff --git a/lib/spack/spack/test/cmd/checksum.py b/lib/spack/spack/test/cmd/checksum.py
new file mode 100644
index 0000000000..6c05b03d59
--- /dev/null
+++ b/lib/spack/spack/test/cmd/checksum.py
@@ -0,0 +1,60 @@
+# Copyright 2013-2021 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)
+
+import argparse
+
+import pytest
+
+import llnl.util.tty as tty
+
+import spack.cmd.checksum
+import spack.repo
+from spack.main import SpackCommand
+
+spack_checksum = SpackCommand('checksum')
+
+
+@pytest.mark.parametrize('arguments,expected', [
+    (['--batch', 'patch'], (True, False, False)),
+    (['--latest', 'patch'], (False, True, False)),
+    (['--preferred', 'patch'], (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
+    assert check == expected
+
+
+@pytest.mark.parametrize('arguments,expected', [
+    (['--batch', 'preferred-test'], 'versions of preferred-test'),
+    (['--latest', 'preferred-test'], 'Found 1 version'),
+    (['--preferred', 'preferred-test'], 'Found 1 version'),
+])
+def test_checksum(arguments, expected, mock_packages, mock_stage):
+    output = spack_checksum(*arguments)
+    assert expected in output
+    assert 'version(' in output
+
+
+def test_checksum_interactive(
+        mock_packages, mock_fetch, mock_stage, monkeypatch):
+    def _get_number(*args, **kwargs):
+        return 1
+    monkeypatch.setattr(tty, 'get_number', _get_number)
+
+    output = spack_checksum('preferred-test')
+    assert 'versions of preferred-test' in output
+    assert 'version(' in output
+
+
+def test_checksum_versions(mock_packages, mock_fetch, mock_stage):
+    pkg = spack.repo.get('preferred-test')
+
+    versions = [str(v) for v in pkg.versions if not v.isdevelop()]
+    output = spack_checksum('preferred-test', versions[0])
+    assert 'Found 1 version' in output
+    assert 'version(' in output
diff --git a/share/spack/spack-completion.bash b/share/spack/spack-completion.bash
index c64495e1d6..bf5300ea96 100755
--- a/share/spack/spack-completion.bash
+++ b/share/spack/spack-completion.bash
@@ -575,7 +575,7 @@ _spack_cd() {
 _spack_checksum() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help --keep-stage -b --batch"
+        SPACK_COMPREPLY="-h --help --keep-stage -b --batch -l --latest -p --preferred"
     else
         _all_packages
     fi