Close #14969: Improve the handling of exception chaining in contextlib.ExitStack

2 files changed, 31 insertions, 14 deletions

diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index f5232b6..bde2feb 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -225,6 +225,17 @@ class ExitStack(object):
         return self
 
     def __exit__(self, *exc_details):
+        # We manipulate the exception state so it behaves as though
+        # we were actually nesting multiple with statements
+        frame_exc = sys.exc_info()[1]
+        def _fix_exception_context(new_exc, old_exc):
+            while 1:
+                exc_context = new_exc.__context__
+                if exc_context in (None, frame_exc):
+                    break
+                new_exc = exc_context
+            new_exc.__context__ = old_exc
+
         # Callbacks are invoked in LIFO order to match the behaviour of
         # nested context managers
         suppressed_exc = False
@@ -236,9 +247,8 @@ class ExitStack(object):
                     exc_details = (None, None, None)
             except:
                 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]
+                # simulate the stack of exceptions by setting the context
+                _fix_exception_context(new_exc_details[1], exc_details[1])
                 if not self._exit_callbacks:
                     raise
                 exc_details = new_exc_details
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index efa9dcb..e52ed91 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -505,6 +505,18 @@ class TestExitStack(unittest.TestCase):
             def __exit__(self, *exc_details):
                 raise self.exc
 
+        class RaiseExcWithContext:
+            def __init__(self, outer, inner):
+                self.outer = outer
+                self.inner = inner
+            def __enter__(self):
+                return self
+            def __exit__(self, *exc_details):
+                try:
+                    raise self.inner
+                except:
+                    raise self.outer
+
         class SuppressExc:
             def __enter__(self):
                 return self
@@ -514,11 +526,10 @@ class TestExitStack(unittest.TestCase):
 
         try:
             with RaiseExc(IndexError):
-                with RaiseExc(KeyError):
-                    with RaiseExc(AttributeError):
-                        with SuppressExc():
-                            with RaiseExc(ValueError):
-                                1 / 0
+                with RaiseExcWithContext(KeyError, AttributeError):
+                    with SuppressExc():
+                        with RaiseExc(ValueError):
+                            1 / 0
         except IndexError as exc:
             self.assertIsInstance(exc.__context__, KeyError)
             self.assertIsInstance(exc.__context__.__context__, AttributeError)
@@ -553,12 +564,8 @@ class TestExitStack(unittest.TestCase):
         except IndexError as exc:
             self.assertIsInstance(exc.__context__, KeyError)
             self.assertIsInstance(exc.__context__.__context__, AttributeError)
-            # Inner exceptions were suppressed, but the with statement
-            # cleanup code adds the one from the body back in as the
-            # context of the exception raised by the outer callbacks
-            # See http://bugs.python.org/issue14969
-            suite_exc = exc.__context__.__context__.__context__
-            self.assertIsInstance(suite_exc, ZeroDivisionError)
+            # Inner exceptions were suppressed
+            self.assertIsNone(exc.__context__.__context__.__context__)
         else:
             self.fail("Expected IndexError, but no exception was raised")
         # Check the inner exceptions