Fix exception context in @contextmanagers

Previously, when using a `try...yield...except` construct within a
`@contextmanager` function, an exception raised by the `yield` wouldn't
be properly cleared when handled the `except` block. Instead, calling
`sys.exception()` would continue to return the exception, leading to
confusing tracebacks and logs.

This was happening due to the way that the `@contextmanager` decorator
drives its decorated generator function. When an exception occurs, it
uses `.throw(exc)` to throw the exception into the generator. However,
even if the generator handles this exception, because the exception was
thrown into it in the context of the `@contextmanager` decorator
handling it (and not being finished yet), `sys.exception()` was not
being reset.

In order to fix this, the exception context as the `@contextmanager`
decorator is `__enter__`'d is stored and set as the current exception
context just before throwing the new exception into the generator.
Doing this means that after the generator handles the thrown exception,
`sys.exception()` reverts back to what it was when the `@contextmanager`
function was started.

Author: pR0Ps

Lib/contextlib.py

@@ -6,6 +6,7 @@
 from collections import deque
 from functools import wraps
 from types import MethodType, GenericAlias
+import ctypes
 
 __all__ = ["asynccontextmanager", "contextmanager", "closing", "nullcontext",
            "AbstractContextManager", "AbstractAsyncContextManager",
@@ -14,6 +15,14 @@
            "chdir"]
 
 
+_PyErr_SetHandledException = ctypes.pythonapi.PyErr_SetHandledException
+_PyErr_SetHandledException.argtypes = [ctypes.py_object]
+def _set_current_exception(exc):
+    if exc is not None and not isinstance(exc, BaseException):
+        raise TypeError(f"{exc!r} is not an exception")
+    _PyErr_SetHandledException(exc)
+
+
 class AbstractContextManager(abc.ABC):
 
     """An abstract base class for context managers."""
@@ -106,6 +115,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
@@ -134,6 +144,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
@@ -152,13 +164,20 @@ def __exit__(self, typ, value, traceback):
                 try:
                     raise RuntimeError("generator didn't stop")
                 finally:
+                    del self.exc_context
                     self.gen.close()
         else:
             if value is None:
                 # Need to force instantiation so we can reliably
                 # tell if we get the same exception back
                 value = typ()
             try:
+                # If the generator handles the exception thrown into it, the
+                # exception context will revert to the actual current exception
+                # context here. In order to make the context manager behave
+                # like a normal function we set the current exception context
+                # to what it was during the context manager's __enter__
+                _set_current_exception(self.exc_context)
                 self.gen.throw(value)
             except StopIteration as exc:
                 # Suppress StopIteration *unless* it's the same exception that
@@ -197,6 +216,7 @@ def __exit__(self, typ, value, traceback):
             try:
                 raise RuntimeError("generator didn't stop after throw()")
             finally:
+                del self.exc_context
                 self.gen.close()
 
 class _AsyncGeneratorContextManager(

Lib/test/test_contextlib.py

@@ -306,6 +306,86 @@ 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_while_handling(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():
+                    try:
+                        raise ValueError()
+                    except:
+                        self.assertIsInstance(sys.exception(), ValueError)
+                        1/0
+                self.assertIsInstance(sys.exception(), TypeError)
+
+        try:
+            raise IndexError()
+        except:
+            try:
+                raise TypeError()
+            except:
+                with self.assertRaises(ZeroDivisionError):
+                    with ctx(reraise=True):
+                        try:
+                            raise ValueError()
+                        except:
+                            self.assertIsInstance(sys.exception(), ValueError)
+                            1/0
+                    self.assertIsInstance(sys.exception(), TypeError)
+
     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.