From 872b1e537e96d0dc4ff37c738031940b5d271366 Mon Sep 17 00:00:00 2001
From: John Belmonte <john@neggie.net>
Date: Tue, 5 Oct 2021 15:21:34 +0900
Subject: [3.10] bpo-44594: fix (Async)ExitStack handling of __context__
 (gh-27089) (GH-28730)

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..
(cherry picked from commit e6d1aa1ac65b6908fdea2c70ec3aa8c4f1dffcb5)

Co-authored-by: John Belmonte <john@neggie.net>

Automerge-Triggered-By: GH:njsmith
---
 Lib/test/test_contextlib.py       | 34 ++++++++++++++++++++++++++++++++++
 Lib/test/test_contextlib_async.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)

(limited to 'Lib/test')

diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index 5a080654ac..fbaae2dd09 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -780,6 +780,40 @@ class TestBaseExitStack:
         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
index 603162eaea..127d750065 100644
--- a/Lib/test/test_contextlib_async.py
+++ b/Lib/test/test_contextlib_async.py
@@ -552,6 +552,41 @@ class TestAsyncExitStack(TestBaseExitStack, unittest.TestCase):
         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")
+
 
 class TestAsyncNullcontext(unittest.TestCase):
     @_async_test
-- 
cgit v1.2.1