Reusable --use-buildcache with better validation (#33388)

Co-authored-by: Massimiliano Culpo <massimiliano.culpo@gmail.com>

4 files changed, 76 insertions, 31 deletions

diff --git a/lib/spack/spack/cmd/common/arguments.py b/lib/spack/spack/cmd/common/arguments.py
index c9f15cfa98..7e68ac594b 100644
--- a/lib/spack/spack/cmd/common/arguments.py
+++ b/lib/spack/spack/cmd/common/arguments.py
@@ -6,6 +6,8 @@
 
 import argparse
 
+from llnl.util.lang import stable_partition
+
 import spack.cmd
 import spack.config
 import spack.dependency as dep
@@ -437,3 +439,57 @@ def add_s3_connection_args(subparser, add_help):
     subparser.add_argument(
         "--s3-endpoint-url", help="Endpoint URL to use to connect to this S3 mirror"
     )
+
+
+def use_buildcache(cli_arg_value):
+    """Translate buildcache related command line arguments into a pair of strings,
+    representing whether the root or its dependencies can use buildcaches.
+
+    Argument type that accepts comma-separated subargs:
+
+        1. auto|only|never
+        2. package:auto|only|never
+        3. dependencies:auto|only|never
+
+    Args:
+        cli_arg_value (str): command line argument value to be translated
+
+    Return:
+        Tuple of two strings
+    """
+    valid_keys = frozenset(["package", "dependencies"])
+    valid_values = frozenset(["only", "never", "auto"])
+
+    # Split in args, split in key/value, and trim whitespace
+    args = [tuple(map(lambda x: x.strip(), part.split(":"))) for part in cli_arg_value.split(",")]
+
+    # Verify keys and values
+    def is_valid(arg):
+        if len(arg) == 1:
+            return arg[0] in valid_values
+        if len(arg) == 2:
+            return arg[0] in valid_keys and arg[1] in valid_values
+        return False
+
+    valid, invalid = stable_partition(args, is_valid)
+
+    # print first error
+    if invalid:
+        raise argparse.ArgumentTypeError("invalid argument `{}`".format(":".join(invalid[0])))
+
+    # Default values
+    package = "auto"
+    dependencies = "auto"
+
+    # Override in order.
+    for arg in valid:
+        if len(arg) == 1:
+            package = dependencies = arg[0]
+            continue
+        key, val = arg
+        if key == "package":
+            package = val
+        else:
+            dependencies = val
+
+    return package, dependencies
diff --git a/lib/spack/spack/cmd/install.py b/lib/spack/spack/cmd/install.py
index a6fdfd31f4..640b48a6d0 100644
--- a/lib/spack/spack/cmd/install.py
+++ b/lib/spack/spack/cmd/install.py
@@ -5,7 +5,6 @@
 
 import argparse
 import os
-import re
 import shutil
 import sys
 import textwrap
@@ -32,33 +31,6 @@ section = "build"
 level = "short"
 
 
-# Pass in the value string passed to use-buildcache and get back
-# the package and dependencies values.
-def parse_use_buildcache(opt):
-    bc_keys = ["package:", "dependencies:", ""]
-    bc_values = ["only", "never", "auto"]
-    kv_list = re.findall("([a-z]+:)?([a-z]+)", opt)
-
-    # Verify keys and values
-    bc_map = {k: v for k, v in kv_list if k in bc_keys and v in bc_values}
-    if not len(kv_list) == len(bc_map):
-        tty.error("Unrecognized arguments passed to use-buildcache")
-        tty.error(
-            "Expected: --use-buildcache "
-            "[[auto|only|never],[package:[auto|only|never]],[dependencies:[auto|only|never]]]"
-        )
-        exit(1)
-
-    for _group in ["package:", "dependencies:"]:
-        if _group not in bc_map:
-            if "" in bc_map:
-                bc_map[_group] = bc_map[""]
-            else:
-                bc_map[_group] = "auto"
-
-    return bc_map["package:"], bc_map["dependencies:"]
-
-
 # Determine value of cache flag
 def cache_opt(default_opt, use_buildcache):
     if use_buildcache == "auto":
@@ -73,8 +45,7 @@ def install_kwargs_from_args(args):
     """Translate command line arguments into a dictionary that will be passed
     to the package installer.
     """
-
-    pkg_use_bc, dep_use_bc = parse_use_buildcache(args.use_buildcache)
+    pkg_use_bc, dep_use_bc = args.use_buildcache
 
     return {
         "fail_fast": args.fail_fast,
@@ -169,6 +140,7 @@ the dependencies""",
     cache_group.add_argument(
         "--use-buildcache",
         dest="use_buildcache",
+        type=arguments.use_buildcache,
         default="package:auto,dependencies:auto",
         metavar="[{auto,only,never},][package:{auto,only,never},][dependencies:{auto,only,never}]",
         help="""select the mode of buildcache for the 'package' and 'dependencies'.
diff --git a/lib/spack/spack/test/cmd/common/arguments.py b/lib/spack/spack/test/cmd/common/arguments.py
index 58bf8e0ac5..772cc0297e 100644
--- a/lib/spack/spack/test/cmd/common/arguments.py
+++ b/lib/spack/spack/test/cmd/common/arguments.py
@@ -129,3 +129,19 @@ def test_concretizer_arguments(mutable_config, mock_packages):
     spec("--fresh", "zlib")
 
     assert spack.config.get("concretizer:reuse", None) is False
+
+
+def test_use_buildcache_type():
+    assert arguments.use_buildcache("only") == ("only", "only")
+    assert arguments.use_buildcache("never") == ("never", "never")
+    assert arguments.use_buildcache("auto") == ("auto", "auto")
+    assert arguments.use_buildcache("package:never,dependencies:only") == ("never", "only")
+    assert arguments.use_buildcache("only,package:never") == ("never", "only")
+    assert arguments.use_buildcache("package:only,package:never") == ("never", "auto")
+    assert arguments.use_buildcache("auto , package: only") == ("only", "auto")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        assert arguments.use_buildcache("pkg:only,deps:never")
+
+    with pytest.raises(argparse.ArgumentTypeError):
+        assert arguments.use_buildcache("sometimes")
diff --git a/lib/spack/spack/test/cmd/install.py b/lib/spack/spack/test/cmd/install.py
index e74f01db7f..8c37ea3e98 100644
--- a/lib/spack/spack/test/cmd/install.py
+++ b/lib/spack/spack/test/cmd/install.py
@@ -17,6 +17,7 @@ from six.moves import builtins
 import llnl.util.filesystem as fs
 import llnl.util.tty as tty
 
+import spack.cmd.common.arguments
 import spack.cmd.install
 import spack.compilers as compilers
 import spack.config
@@ -1132,7 +1133,7 @@ def test_install_use_buildcache(
             "--no-check-signature", "--use-buildcache", opt, package_name, fail_on_error=True
         )
 
-        pkg_opt, dep_opt = spack.cmd.install.parse_use_buildcache(opt)
+        pkg_opt, dep_opt = spack.cmd.common.arguments.use_buildcache(opt)
         validate(dep_opt, out, dependency_name)
         validate(pkg_opt, out, package_name)