gh-111375: Fix handling of exceptions within @contextmanager-decorated functions

Lib/contextlib.py

@@ -168,6 +168,7 @@ class _GeneratorContextManagerBase:
     """Shared functionality for @contextmanager and @asynccontextmanager."""
 
     def __init__(self, func, args, kwds):
+        self.exc_context = None
         self.gen = func(*args, **kwds)
         self.func, self.args, self.kwds = func, args, kwds
         # Issue 19330: ensure context manager instances have good docstrings
@@ -196,6 +197,8 @@ class _GeneratorContextManager(
     """Helper for @contextmanager decorator."""
 
     def __enter__(self):
+        # store the exception context on enter so it can be restored on exit
+        self.exc_context = sys.exception()
         # do not keep args and kwds alive unnecessarily
         # they are only needed for recreation, which is not possible anymore
         del self.args, self.kwds, self.func
@@ -205,6 +208,9 @@ def __enter__(self):
             raise RuntimeError("generator didn't yield") from None
 
     def __exit__(self, typ, value, traceback):
+        # don't keep the stored exception alive unnecessarily
+        exc_context = self.exc_context
+        self.exc_context = None
         if typ is None:
             try:
                 next(self.gen)
@@ -221,7 +227,18 @@ def __exit__(self, typ, value, traceback):
                 # tell if we get the same exception back
                 value = typ()
             try:
-                self.gen.throw(value)
+                # Throw the current exception into the generator so it can
+                # handle it.
+                # Once the generator handles the thrown exception, the
+                # exception context within it should revert back to what it was
+                # before its "yield" statement (ie. what it was in self.__enter__).
+                # However, since we're still currently handling the exception
+                # that we throw into the generator here, the exception context
+                # in the generator wouldn't change.
+                # To work around this, we pass the exception context as it was
+                # just before the generator's yield statement into the generator.
+                # (see gh-111676).
+                self.gen.throw(value, exc_context=exc_context)
             except StopIteration as exc:
                 # Suppress StopIteration *unless* it's the same exception that
                 # was passed to throw().  This prevents a StopIteration

Lib/test/test_contextlib.py

@@ -306,6 +306,98 @@ def woohoo():
             with woohoo():
                 raise StopIteration
 
+    def test_contextmanager_handling_exception_resets_exc_info(self):
+        # Test that sys.exc_info() is correctly unset after handling the error
+        # when used within a context manager
+
+        @contextmanager
+        def ctx(reraise=False):
+            try:
+                self.assertIsNone(sys.exception())
+                yield
+            except:
+                self.assertIsInstance(sys.exception(), ZeroDivisionError)
+                if reraise:
+                    raise
+            else:
+                self.assertIsNone(sys.exception())
+            self.assertIsNone(sys.exception())
+
+        with ctx():
+            pass
+
+        with ctx():
+            1/0
+
+        with self.assertRaises(ZeroDivisionError):
+            with ctx(reraise=True):
+                1/0
+
+    def test_contextmanager_preserves_handled_exception(self):
+        # test that any exceptions currently being handled are preserved
+        # through the context manager
+
+        @contextmanager
+        def ctx(reraise=False):
+            # called while handling an IndexError --> TypeError
+            self.assertIsInstance(sys.exception(), TypeError)
+            self.assertIsInstance(sys.exception().__context__, IndexError)
+            exc_ctx = sys.exception()
+            try:
+                # raises a ValueError --> ZeroDivisionError
+                yield
+            except:
+                self.assertIsInstance(sys.exception(), ZeroDivisionError)
+                self.assertIsInstance(sys.exception().__context__, ValueError)
+                # original error context is preserved
+                self.assertIs(sys.exception().__context__.__context__, exc_ctx)
+                if reraise:
+                    raise
+
+            # inner error handled, context should now be the original context
+            self.assertIs(sys.exception(), exc_ctx)
+
+        try:
+            raise IndexError()
+        except:
+            try:
+                raise TypeError()
+            except:
+                with ctx():
+                    self.assertIsInstance(sys.exception(), TypeError)
+                    self.assertIsInstance(sys.exception().__context__, IndexError)
+                    try:
+                        raise ValueError()
+                    except:
+                        self.assertIsInstance(sys.exception(), ValueError)
+                        self.assertIsInstance(sys.exception().__context__, TypeError)
+                        self.assertIsInstance(sys.exception().__context__.__context__, IndexError)
+                        1/0
+                self.assertIsInstance(sys.exception(), TypeError)
+                self.assertIsInstance(sys.exception().__context__, IndexError)
+            self.assertIsInstance(sys.exception(), IndexError)
+
+        try:
+            raise IndexError()
+        except:
+            try:
+                raise TypeError()
+            except:
+                with self.assertRaises(ZeroDivisionError):
+                    with ctx(reraise=True):
+                        self.assertIsInstance(sys.exception(), TypeError)
+                        self.assertIsInstance(sys.exception().__context__, IndexError)
+                        try:
+                            raise ValueError()
+                        except:
+                            self.assertIsInstance(sys.exception(), ValueError)
+                            self.assertIsInstance(sys.exception().__context__, TypeError)
+                            self.assertIsInstance(sys.exception().__context__.__context__, IndexError)
+                            1/0
+                    self.assertIsInstance(sys.exception(), TypeError)
+                    self.assertIsInstance(sys.exception().__context__, IndexError)
+            self.assertIsInstance(sys.exception(), IndexError)
+
     def _create_contextmanager_attribs(self):
         def attribs(**kw):
             def decorate(func):

Misc/NEWS.d/next/Library/2023-11-03-01-44-51.gh-issue-111375.qDBcCI.rst

@@ -0,0 +1,2 @@
+Fix handling of ``sys.exception()`` within ``@contextlib.contextmanager``
+functions. Patch by Carey Metcalfe.