From f05837a4801f5f5606cf26f2a7754a1ab1d67495 Mon Sep 17 00:00:00 2001
From: Harmen Stoppels <harmenstoppels@gmail.com>
Date: Mon, 17 Jul 2023 23:20:24 +0200
Subject: [PATCH] Fix wrong StageComposite keep override (#38938)

`Stage(keep=True)` was ignored when put in a composite that doesn't
override the value.
---
 lib/spack/spack/stage.py      | 19 +++++++++++++++++--
 lib/spack/spack/test/stage.py | 21 +++++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/lib/spack/spack/stage.py b/lib/spack/spack/stage.py
index 2dd3a98a554298..6a5297dce2d754 100644
--- a/lib/spack/spack/stage.py
+++ b/lib/spack/spack/stage.py
@@ -12,7 +12,7 @@
 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 @@ def __init__(self):
             ]
         )
 
+    @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 @@ def __enter__(self):
 
     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 @@ def path(self):
     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 1220a86904ad9c..e507ac74e363cb 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