Add more functionality to the stage cmd (#46498)

* Add more functionality to the stage cmd

* Completion commands

* completion again

* Add tests, but they are slow

* Stale comment

lib/spack/spack/cmd/stage.py

@@ -19,11 +19,48 @@
 level = "long"
 
 
+class StageFilter:
+    """
+    Encapsulation of reasons to skip staging
+    """
+
+    def __init__(self, exclusions, skip_installed):
+        """
+        :param exclusions: A list of specs to skip if satisfied.
+        :param skip_installed: A boolean indicating whether to skip already installed specs.
+        """
+        self.exclusions = exclusions
+        self.skip_installed = skip_installed
+
+    def __call__(self, spec):
+        """filter action, true means spec should be filtered"""
+        if spec.external:
+            return True
+
+        if self.skip_installed and spec.installed:
+            return True
+
+        if any(spec.satisfies(exclude) for exclude in self.exclusions):
+            return True
+
+        return False
+
+
 def setup_parser(subparser):
     arguments.add_common_arguments(subparser, ["no_checksum", "specs"])
     subparser.add_argument(
         "-p", "--path", dest="path", help="path to stage package, does not add to spack tree"
     )
+    subparser.add_argument(
+        "-e",
+        "--exclude",
+        action="append",
+        default=[],
+        help="exclude packages that satisfy the specified specs",
+    )
+    subparser.add_argument(
+        "-s", "--skip-installed", action="store_true", help="dont restage already installed specs"
+    )
     arguments.add_concretizer_args(subparser)
 
 
@@ -31,11 +68,14 @@ def stage(parser, args):
     if args.no_checksum:
         spack.config.set("config:checksum", False, scope="command_line")
 
+    exclusion_specs = spack.cmd.parse_specs(args.exclude, concretize=False)
+    filter = StageFilter(exclusion_specs, args.skip_installed)
+
     if not args.specs:
         env = ev.active_environment()
         if not env:
             tty.die("`spack stage` requires a spec or an active environment")
-        return _stage_env(env)
+        return _stage_env(env, filter)
 
     specs = spack.cmd.parse_specs(args.specs, concretize=False)
 
@@ -49,6 +89,11 @@ def stage(parser, args):
 
     specs = spack.cmd.matching_specs_from_env(specs)
     for spec in specs:
+        spec = spack.cmd.matching_spec_from_env(spec)
+
+        if filter(spec):
+            continue
+
         pkg = spec.package
 
         if custom_path:
@@ -57,9 +102,13 @@ def stage(parser, args):
         _stage(pkg)
 
 
-def _stage_env(env: ev.Environment):
+def _stage_env(env: ev.Environment, filter):
     tty.msg(f"Staging specs from environment {env.name}")
     for spec in spack.traverse.traverse_nodes(env.concrete_roots()):
+
+        if filter(spec):
+            continue
+
         _stage(spec.package)
 
 

lib/spack/spack/test/cmd/stage.py

@@ -11,7 +11,9 @@
 import spack.environment as ev
 import spack.package_base
 import spack.traverse
+from spack.cmd.stage import StageFilter
 from spack.main import SpackCommand, SpackCommandError
+from spack.spec import Spec
 from spack.version import Version
 
 stage = SpackCommand("stage")
@@ -127,3 +129,54 @@ def test_concretizer_arguments(mock_packages, mock_fetch):
 
     stage("--fresh", "trivial-install-test-package")
     assert spack.config.get("concretizer:reuse", None) is False
+
+
+@pytest.mark.maybeslow
+@pytest.mark.parametrize("externals", [["libelf"], []])
+@pytest.mark.parametrize(
+    "installed, skip_installed", [(["libdwarf"], False), (["libdwarf"], True)]
+)
+@pytest.mark.parametrize("exclusions", [["mpich", "callpath"], []])
+def test_stage_spec_filters(
+    mutable_mock_env_path,
+    mock_packages,
+    mock_fetch,
+    externals,
+    installed,
+    skip_installed,
+    exclusions,
+    monkeypatch,
+):
+    e = ev.create("test")
+    e.add("mpileaks@=100.100")
+    e.concretize()
+    all_specs = e.all_specs()
+
+    def is_installed(self):
+        return self.name in installed
+
+    if skip_installed:
+        monkeypatch.setattr(Spec, "installed", is_installed)
+
+    should_be_filtered = []
+    for spec in all_specs:
+        for ext in externals:
+            if spec.satisfies(Spec(ext)):
+                spec.external_path = "/usr"
+                assert spec.external
+                should_be_filtered.append(spec)
+        for ins in installed:
+            if skip_installed and spec.satisfies(Spec(ins)):
+                assert spec.installed
+                should_be_filtered.append(spec)
+        for exc in exclusions:
+            if spec.satisfies(Spec(exc)):
+                should_be_filtered.append(spec)
+
+    filter = StageFilter(exclusions, skip_installed=skip_installed)
+    specs_to_stage = [s for s in all_specs if not filter(s)]
+    specs_were_filtered = [skip not in specs_to_stage for skip in should_be_filtered]
+
+    assert all(
+        specs_were_filtered
+    ), f"Packages associated with bools: {[s.name for s in should_be_filtered]}"