spack checksum: restore ability to select top n (#40531)

The ability to select the top N versions got removed in the checksum overhaul,
cause initially numbers were used for commands.

Now that we settled on characters for commands, let's make numbers pick the top
N again.

2 files changed, 40 insertions, 6 deletions

diff --git a/lib/spack/spack/stage.py b/lib/spack/spack/stage.py
index d53ec5fee8..1c7ebdec5c 100644
--- a/lib/spack/spack/stage.py
+++ b/lib/spack/spack/stage.py
@@ -921,7 +921,7 @@ def interactive_version_filter(
 
         print_header = True
 
-        print("commands:")
+        tty.info(colorize("Enter @*{number} of versions to take, or use a @*{command}:"))
         commands = (
             "@*b{[c]}hecksum",
             "@*b{[e]}dit",
@@ -931,10 +931,10 @@ def interactive_version_filter(
             "@*b{[r]}estart",
             "@*b{[q]}uit",
         )
-        colify(list(map(colorize, commands)), indent=2)
+        colify(list(map(colorize, commands)), indent=4)
 
         try:
-            command = input(colorize("@*g{command>} ")).strip().lower()
+            command = input(colorize("@*g{action>} ")).strip().lower()
         except EOFError:
             print()
             command = "q"
@@ -1039,9 +1039,20 @@ def interactive_version_filter(
                 print()
                 return None
         else:
-            tty.warn(f"Ignoring invalid command: {command}")
-            print_header = False
-            continue
+            # Last restort: filter the top N versions
+            try:
+                n = int(command)
+                invalid_command = n < 1
+            except ValueError:
+                invalid_command = True
+
+            if invalid_command:
+                tty.warn(f"Ignoring invalid command: {command}")
+                print_header = False
+                continue
+
+            sorted_and_filtered = sorted_and_filtered[:n]
+
     return {v: url_dict[v] for v in sorted_and_filtered}
 
 
diff --git a/lib/spack/spack/test/cmd/checksum.py b/lib/spack/spack/test/cmd/checksum.py
index 8001334e3e..b2fc9d5f6c 100644
--- a/lib/spack/spack/test/cmd/checksum.py
+++ b/lib/spack/spack/test/cmd/checksum.py
@@ -202,6 +202,29 @@ def test_checksum_interactive_new_only():
     }
 
 
+def test_checksum_interactive_top_n():
+    """Test integers select top n versions"""
+    input = input_from_commands("2", "c")
+    assert interactive_version_filter(
+        {
+            Version("1.1"): "https://www.example.com/pkg-1.1.tar.gz",
+            Version("1.0"): "https://www.example.com/pkg-1.0.tar.gz",
+            Version("0.9"): "https://www.example.com/pkg-0.9.tar.gz",
+        },
+        input=input,
+    ) == {
+        Version("1.1"): "https://www.example.com/pkg-1.1.tar.gz",
+        Version("1.0"): "https://www.example.com/pkg-1.0.tar.gz",
+    }
+
+
+def test_checksum_interactive_unrecognized_command():
+    """Unrecognized commands should be ignored"""
+    input = input_from_commands("-1", "0", "hello", "c")
+    v = {Version("1.1"): "https://www.example.com/pkg-1.1.tar.gz"}
+    assert interactive_version_filter(v.copy(), input=input) == v
+
+
 def test_checksum_versions(mock_packages, mock_clone_repo, mock_fetch, mock_stage):
     pkg_cls = spack.repo.PATH.get_pkg_class("zlib")
     versions = [str(v) for v in pkg_cls.versions]