bpo-44594: fix (Async)ExitStack handling of __context__ (gh-27089)

belm0 · web-flow · commit e6d1aa1ac65b · 2021-10-03T23:49:55.000-07:00
* bpo-44594: fix (Async)ExitStack handling of __context__ Make enter_context(foo()) / enter_async_context(foo()) equivalent to `[async] with foo()` regarding __context__ when an exception is raised. Previously exceptions would be caught and re-raised with the wrong context when explicitly overriding __context__ with None.
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
@@ -553,10 +553,10 @@ def _fix_exception_context(new_exc, old_exc):
             # Context may not be correct, so find the end of the chain
             while 1:
                 exc_context = new_exc.__context__
-                if exc_context is old_exc:
+                if exc_context is None or exc_context is old_exc:
                     # Context is already set correctly (see issue 20317)
                     return
-                if exc_context is None or exc_context is frame_exc:
+                if exc_context is frame_exc:
                     break
                 new_exc = exc_context
             # Change the end of the chain to point to the exception
@@ -693,10 +693,10 @@ def _fix_exception_context(new_exc, old_exc):
             # Context may not be correct, so find the end of the chain
             while 1:
                 exc_context = new_exc.__context__
-                if exc_context is old_exc:
+                if exc_context is None or exc_context is old_exc:
                     # Context is already set correctly (see issue 20317)
                     return
-                if exc_context is None or exc_context is frame_exc:
+                if exc_context is frame_exc:
                     break
                 new_exc = exc_context
             # Change the end of the chain to point to the exception
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
@@ -799,6 +799,40 @@ def suppress_exc(*exc_details):
         self.assertIsInstance(inner_exc, ValueError)
         self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
 
+    def test_exit_exception_explicit_none_context(self):
+        # Ensure ExitStack chaining matches actual nested `with` statements
+        # regarding explicit __context__ = None.
+
+        class MyException(Exception):
+            pass
+
+        @contextmanager
+        def my_cm():
+            try:
+                yield
+            except BaseException:
+                exc = MyException()
+                try:
+                    raise exc
+                finally:
+                    exc.__context__ = None
+
+        @contextmanager
+        def my_cm_with_exit_stack():
+            with self.exit_stack() as stack:
+                stack.enter_context(my_cm())
+                yield stack
+
+        for cm in (my_cm, my_cm_with_exit_stack):
+            with self.subTest():
+                try:
+                    with cm():
+                        raise IndexError()
+                except MyException as exc:
+                    self.assertIsNone(exc.__context__)
+                else:
+                    self.fail("Expected IndexError, but no exception was raised")
+
     def test_exit_exception_non_suppressing(self):
         # http://bugs.python.org/issue19092
         def raise_exc(exc):
diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py
@@ -646,6 +646,41 @@ async def suppress_exc(*exc_details):
         self.assertIsInstance(inner_exc, ValueError)
         self.assertIsInstance(inner_exc.__context__, ZeroDivisionError)
 
+    @_async_test
+    async def test_async_exit_exception_explicit_none_context(self):
+        # Ensure AsyncExitStack chaining matches actual nested `with` statements
+        # regarding explicit __context__ = None.
+
+        class MyException(Exception):
+            pass
+
+        @asynccontextmanager
+        async def my_cm():
+            try:
+                yield
+            except BaseException:
+                exc = MyException()
+                try:
+                    raise exc
+                finally:
+                    exc.__context__ = None
+
+        @asynccontextmanager
+        async def my_cm_with_exit_stack():
+            async with self.exit_stack() as stack:
+                await stack.enter_async_context(my_cm())
+                yield stack
+
+        for cm in (my_cm, my_cm_with_exit_stack):
+            with self.subTest():
+                try:
+                    async with cm():
+                        raise IndexError()
+                except MyException as exc:
+                    self.assertIsNone(exc.__context__)
+                else:
+                    self.fail("Expected IndexError, but no exception was raised")
+
     @_async_test
     async def test_instance_bypass_async(self):
         class Example(object): pass
diff --git a/Misc/NEWS.d/next/Library/2021-07-12-10-32-48.bpo-44594.eEa5zi.rst b/Misc/NEWS.d/next/Library/2021-07-12-10-32-48.bpo-44594.eEa5zi.rst
@@ -0,0 +1,3 @@
+Fix an edge case of :class:`ExitStack` and :class:`AsyncExitStack` exception
+chaining.  They will now match ``with`` block behavior when ``__context__`` is
+explicitly set to ``None`` when the exception is in flight.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+Fix an edge case of :class:`ExitStack` and :class:`AsyncExitStack` exception
	`2`	+chaining. They will now match ``with`` block behavior when ``__context__`` is
	`3`	+explicitly set to ``None`` when the exception is in flight.