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

Lib/contextlib.py

@@ -20,6 +20,12 @@
            "chdir"]
 
 
+try:
+    from _contextlib import _setup_genmgr_ctx
+except ImportError:
+    _setup_genmgr_ctx = None
+
+
 class AbstractContextManager(abc.ABC):
 
     """An abstract base class for context managers."""
@@ -168,6 +174,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 +203,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 +214,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,6 +233,20 @@ def __exit__(self, typ, value, traceback):
                 # tell if we get the same exception back
                 value = typ()
             try:
+                # 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 forcefully set the current exception
+                # context to be what it was just before the generator's yield
+                # statement before throwing the current exception into it.
+                # (see gh-111676).
+                if _setup_genmgr_ctx is not None:
+                    _setup_genmgr_ctx(exc_context)
                 self.gen.throw(value)
             except StopIteration as exc:
                 # Suppress StopIteration *unless* it's the same exception that

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):
@@ -1579,5 +1671,138 @@ def test_exception(self):
         self.assertEqual(os.getcwd(), old_cwd)
 
 
+try:
+    from _contextlib import _setup_genmgr_ctx
+except ImportError:
+    _setup_genmgr_ctx = None
+
+if _setup_genmgr_ctx:
+    class SetExceptionTests(unittest.TestCase):
+
+        def tearDown(self):
+            # make sure we don't leave the global exception set
+            _setup_genmgr_ctx(None);
+
+        def test_set_exc_invalid_values(self):
+            for x in (0, "1", b"2"):
+                with self.assertRaises(TypeError):
+                    _setup_genmgr_ctx(x);
+
+        def test_clear_exc(self):
+            try:
+                raise ValueError()
+            except ValueError:
+                self.assertIsInstance(sys.exception(), ValueError)
+                _setup_genmgr_ctx(None)
+                self.assertIsNone(sys.exception())
+
+        def test_set_exc(self):
+            exc = ValueError()
+            self.assertIsNone(sys.exception())
+            _setup_genmgr_ctx(exc)
+            self.assertIs(sys.exception(), exc)
+
+        def test_set_exc_replaced_by_new_exception_and_restored(self):
+            exc = ValueError()
+            _setup_genmgr_ctx(exc)
+            self.assertIs(sys.exception(), exc)
+            try:
+                raise TypeError()
+            except TypeError:
+                self.assertIsInstance(sys.exception(), TypeError)
+                self.assertIs(sys.exception().__context__, exc)
+
+            self.assertIs(sys.exception(), exc)
+
+        def test_set_exc_popped_on_exit_except(self):
+            exc = ValueError()
+            try:
+                raise TypeError()
+            except TypeError:
+                self.assertIsInstance(sys.exception(), TypeError)
+                _setup_genmgr_ctx(exc)
+                self.assertIs(sys.exception(), exc)
+            self.assertIsNone(sys.exception())
+
+        def test_cleared_exc_overridden_and_restored(self):
+            try:
+                raise ValueError()
+            except ValueError:
+                try:
+                    raise TypeError()
+                except TypeError:
+                    self.assertIsInstance(sys.exception(), TypeError)
+                    _setup_genmgr_ctx(None)
+                    self.assertIsNone(sys.exception())
+                    try:
+                        raise IndexError()
+                    except IndexError:
+                        self.assertIsInstance(sys.exception(), IndexError)
+                        self.assertIsNone(sys.exception().__context__)
+                    self.assertIsNone(sys.exception())
+                self.assertIsInstance(sys.exception(), ValueError)
+            self.assertIsNone(sys.exception())
+
+        def test_clear_exc_in_generator(self):
+            def inner():
+                self.assertIsNone(sys.exception())
+                yield
+                self.assertIsInstance(sys.exception(), ValueError)
+                _setup_genmgr_ctx(None)
+                self.assertIsNone(sys.exception())
+                yield
+                self.assertIsNone(sys.exception())
+
+            # with a single exception in exc_info stack
+            g = inner()
+            next(g)
+            try:
+                raise ValueError()
+            except:
+                self.assertIsInstance(sys.exception(), ValueError)
+                next(g)
+                self.assertIsInstance(sys.exception(), ValueError)
+            self.assertIsNone(sys.exception())
+            with self.assertRaises(StopIteration):
+                next(g)
+            self.assertIsNone(sys.exception())
+
+            # with multiple exceptions in exc_info stack by chaining generators
+            def outer():
+                g = inner()
+                self.assertIsNone(sys.exception())
+                yield next(g)
+                self.assertIsInstance(sys.exception(), TypeError)
+                try:
+                    raise ValueError()
+                except:
+                    self.assertIsInstance(sys.exception(), ValueError)
+                    self.assertIsInstance(sys.exception().__context__, TypeError)
+                    yield next(g)
+                    # at this point the TypeError from the caller has been handled
+                    # by the caller's except block. Even still, it should still be
+                    # referenced as the __context__ of the current exception.
+                    self.assertIsInstance(sys.exception(), ValueError)
+                    self.assertIsInstance(sys.exception().__context__, TypeError)
+                # not handling an exception, caller isn't handling one either
+                self.assertIsNone(sys.exception())
+                with self.assertRaises(StopIteration):
+                    next(g)
+                self.assertIsNone(sys.exception())
+
+            g = outer()
+            next(g)
+            try:
+                raise TypeError()
+            except:
+                self.assertIsInstance(sys.exception(), TypeError)
+                next(g)
+                self.assertIsInstance(sys.exception(), TypeError)
+            self.assertIsNone(sys.exception())
+            with self.assertRaises(StopIteration):
+                next(g)
+            self.assertIsNone(sys.exception())
+
+
 if __name__ == "__main__":
     unittest.main()

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.

Modules/Setup.bootstrap.in

@@ -31,6 +31,7 @@ _weakref _weakref.c
 
 # commonly used core modules
 _abc _abc.c
+_contextlib _contextlibmodule.c
 _functools _functoolsmodule.c
 _locale _localemodule.c
 _opcode _opcode.c

Modules/_contextlibmodule.c

@@ -0,0 +1,72 @@
+
+/* Contextlib module */
+
+#include "Python.h"
+
+#include "clinic/_contextlibmodule.c.h"
+
+/*[clinic input]
+module _contextlib
+[clinic start generated code]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=60c4760e25c64b5e]*/
+
+
+/*[clinic input]
+_contextlib._setup_genmgr_ctx
+    exception: object
+
+Force the provided exception to be set as the current exception
+
+Subsequent calls to sys.exception()/sys.exc_info() will return
+the provided exception until another exception is caught in the
+current thread or the execution stack returns to a frame where
+another exception is being handled.
+[clinic start generated code]*/
+
+static PyObject *
+_contextlib__setup_genmgr_ctx_impl(PyObject *module, PyObject *exception)
+/*[clinic end generated code: output=d8996a9f8f52204b input=d61410ba5659bf6f]*/
+{
+    if (!Py_IsNone(exception) && !PyExceptionInstance_Check(exception)){
+        PyErr_SetString(
+            PyExc_TypeError,
+            "must be an exception/None"
+        );
+        return NULL;
+    }
+    PyErr_SetHandledException(exception);
+    Py_RETURN_NONE;
+}
+
+
+/* module level code ********************************************************/
+
+PyDoc_STRVAR(_contextlib_doc,
+"Helpers for contextlib.");
+
+static PyMethodDef contextlib_methods[] = {
+    _CONTEXTLIB__SETUP_GENMGR_CTX_METHODDEF
+    {NULL,              NULL}           /* sentinel */
+};
+
+static PyModuleDef_Slot contextlib_slots[] = {
+    _Py_ABI_SLOT,
+    {Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},
+    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
+    {0, NULL}
+};
+
+static struct PyModuleDef _contextlibmodule = {
+    PyModuleDef_HEAD_INIT,
+    .m_name = "_contextlib",
+    .m_doc = _contextlib_doc,
+    .m_size = 0,
+    .m_methods = contextlib_methods,
+    .m_slots = contextlib_slots,
+};
+
+PyMODINIT_FUNC
+PyInit__contextlib(void)
+{
+    return PyModuleDef_Init(&_contextlibmodule);
+}