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]}"

share/spack/spack-completion.bash

@@ -1867,7 +1867,7 @@ _spack_spec() {
 _spack_stage() {
     if $list_options
     then
-        SPACK_COMPREPLY="-h --help -n --no-checksum -p --path -U --fresh --reuse --fresh-roots --reuse-deps --deprecated"
+        SPACK_COMPREPLY="-h --help -n --no-checksum -p --path -e --exclude -s --skip-installed -U --fresh --reuse --fresh-roots --reuse-deps --deprecated"
     else
         _all_packages
     fi

share/spack/spack-completion.fish

@@ -2885,7 +2885,7 @@ complete -c spack -n '__fish_spack_using_command spec' -l deprecated -f -a confi
 complete -c spack -n '__fish_spack_using_command spec' -l deprecated -d 'allow concretizer to select deprecated versions'
 
 # spack stage
-set -g __fish_spack_optspecs_spack_stage h/help n/no-checksum p/path= U/fresh reuse fresh-roots deprecated
+set -g __fish_spack_optspecs_spack_stage h/help n/no-checksum p/path= e/exclude= s/skip-installed U/fresh reuse fresh-roots deprecated
 complete -c spack -n '__fish_spack_using_command_pos_remainder 0 stage' -f -k -a '(__fish_spack_specs_or_id)'
 complete -c spack -n '__fish_spack_using_command stage' -s h -l help -f -a help
 complete -c spack -n '__fish_spack_using_command stage' -s h -l help -d 'show this help message and exit'
@@ -2893,6 +2893,10 @@ complete -c spack -n '__fish_spack_using_command stage' -s n -l no-checksum -f -
 complete -c spack -n '__fish_spack_using_command stage' -s n -l no-checksum -d 'do not use checksums to verify downloaded files (unsafe)'
 complete -c spack -n '__fish_spack_using_command stage' -s p -l path -r -f -a path
 complete -c spack -n '__fish_spack_using_command stage' -s p -l path -r -d 'path to stage package, does not add to spack tree'
+complete -c spack -n '__fish_spack_using_command stage' -s e -l exclude -r -f -a exclude
+complete -c spack -n '__fish_spack_using_command stage' -s e -l exclude -r -d 'exclude packages that satisfy the specified specs'
+complete -c spack -n '__fish_spack_using_command stage' -s s -l skip-installed -f -a skip_installed
+complete -c spack -n '__fish_spack_using_command stage' -s s -l skip-installed -d 'dont restage already installed specs'
 complete -c spack -n '__fish_spack_using_command stage' -s U -l fresh -f -a concretizer_reuse
 complete -c spack -n '__fish_spack_using_command stage' -s U -l fresh -d 'do not reuse installed deps; build newest configuration'
 complete -c spack -n '__fish_spack_using_command stage' -l reuse -f -a concretizer_reuse