Fix wrong StageComposite keep override (#38938)

`Stage(keep=True)` was ignored when put in a composite that doesn't
override the value.

2 files changed, 38 insertions, 2 deletions

diff --git a/lib/spack/spack/stage.py b/lib/spack/spack/stage.py
index 2dd3a98a55..6a5297dce2 100644
--- a/lib/spack/spack/stage.py
+++ b/lib/spack/spack/stage.py
@@ -12,7 +12,7 @@ import shutil
 import stat
 import sys
 import tempfile
-from typing import Dict
+from typing import Dict, Iterable
 
 import llnl.util.lang
 import llnl.util.tty as tty
@@ -742,6 +742,13 @@ class StageComposite(pattern.Composite):
             ]
         )
 
+    @classmethod
+    def from_iterable(cls, iterable: Iterable[Stage]) -> "StageComposite":
+        """Create a new composite from an iterable of stages."""
+        composite = cls()
+        composite.extend(iterable)
+        return composite
+
     def __enter__(self):
         for item in self:
             item.__enter__()
@@ -749,7 +756,6 @@ class StageComposite(pattern.Composite):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         for item in reversed(self):
-            item.keep = getattr(self, "keep", False)
             item.__exit__(exc_type, exc_val, exc_tb)
 
     #
@@ -771,6 +777,15 @@ class StageComposite(pattern.Composite):
     def archive_file(self):
         return self[0].archive_file
 
+    @property
+    def keep(self):
+        return self[0].keep
+
+    @keep.setter
+    def keep(self, value):
+        for item in self:
+            item.keep = value
+
 
 class DIYStage:
     """
diff --git a/lib/spack/spack/test/stage.py b/lib/spack/spack/test/stage.py
index 1220a86904..e507ac74e3 100644
--- a/lib/spack/spack/test/stage.py
+++ b/lib/spack/spack/test/stage.py
@@ -890,3 +890,24 @@ def test_cannot_access(capsys):
 
     captured = capsys.readouterr()
     assert "Insufficient permissions" in str(captured)
+
+
+def test_override_keep_in_composite_stage():
+    stage_1 = Stage("file:///does-not-exist", keep=True)
+    stage_2 = Stage("file:///does-not-exist", keep=False)
+    stage_3 = Stage("file:///does-not-exist", keep=True)
+    stages = spack.stage.StageComposite.from_iterable((stage_1, stage_2, stage_3))
+
+    # The getter for the composite stage just returns the value of the first stage
+    # its just there so we have a setter too.
+    assert stages.keep
+    assert stage_1.keep
+    assert not stage_2.keep
+    assert stage_3.keep
+
+    # This should override all stages
+    stages.keep = False
+    assert not stages.keep
+    assert not stage_1.keep
+    assert not stage_2.keep
+    assert not stage_3.keep