Close #14963: Use an iterative algorithm in contextlib.ExitStack.__exit__ (Patch by Alon Horev)

2 files changed, 21 insertions, 26 deletions

diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index ead1155..f5232b6 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -225,32 +225,21 @@ class ExitStack(object):
         return self
 
     def __exit__(self, *exc_details):
-        if not self._exit_callbacks:
-            return
-        # This looks complicated, but it is really just
-        # setting up a chain of try-expect statements to ensure
-        # that outer callbacks still get invoked even if an
-        # inner one throws an exception
-        def _invoke_next_callback(exc_details):
-            # Callbacks are removed from the list in FIFO order
-            # but the recursion means they're invoked in LIFO order
-            cb = self._exit_callbacks.popleft()
-            if not self._exit_callbacks:
-                # Innermost callback is invoked directly
-                return cb(*exc_details)
-            # More callbacks left, so descend another level in the stack
+        # Callbacks are invoked in LIFO order to match the behaviour of
+        # nested context managers
+        suppressed_exc = False
+        while self._exit_callbacks:
+            cb = self._exit_callbacks.pop()
             try:
-                suppress_exc = _invoke_next_callback(exc_details)
+                if cb(*exc_details):
+                    suppressed_exc = True
+                    exc_details = (None, None, None)
             except:
-                suppress_exc = cb(*sys.exc_info())
-                # Check if this cb suppressed the inner exception
-                if not suppress_exc:
+                new_exc_details = sys.exc_info()
+                if exc_details != (None, None, None):
+                    # simulate the stack of exceptions by setting the context
+                    new_exc_details[1].__context__ = exc_details[1]
+                if not self._exit_callbacks:
                     raise
-            else:
-                # Check if inner cb suppressed the original exception
-                if suppress_exc:
-                    exc_details = (None, None, None)
-                suppress_exc = cb(*exc_details) or suppress_exc
-            return suppress_exc
-        # Kick off the recursive chain
-        return _invoke_next_callback(exc_details)
+                exc_details = new_exc_details
+        return suppressed_exc
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index e5eed21..efa9dcb 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -572,6 +572,12 @@ class TestExitStack(unittest.TestCase):
             stack.push(lambda *exc: 1/0)
             stack.push(lambda *exc: {}[1])
 
+    def test_excessive_nesting(self):
+        # The original implementation would die with RecursionError here
+        with ExitStack() as stack:
+            for i in range(10000):
+                stack.callback(int)
+
     def test_instance_bypass(self):
         class Example(object): pass
         cm = Example()