diff --git a/Lib/contextlib.py b/Lib/contextlib.py index efc02bfa9243da6..771432cec4cb636 100644 --- a/Lib/contextlib.py +++ b/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 diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py index e291f814edbd930..73c1bb7e32bafed 100644 --- a/Lib/test/test_contextlib.py +++ b/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): diff --git a/Misc/NEWS.d/next/Library/2023-11-03-01-44-51.gh-issue-111375.qDBcCI.rst b/Misc/NEWS.d/next/Library/2023-11-03-01-44-51.gh-issue-111375.qDBcCI.rst new file mode 100644 index 000000000000000..2519b7e614e9a33 --- /dev/null +++ b/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. diff --git a/Objects/genobject.c b/Objects/genobject.c index 8c5d720c0b9035c..61c778e8177fdcb 100644 --- a/Objects/genobject.c +++ b/Objects/genobject.c @@ -612,8 +612,8 @@ gen_throw_current_exception(PyGenObject *gen) } PyDoc_STRVAR(throw_doc, -"throw(value)\n\ -throw(type[,value[,tb]])\n\ +"throw(value, /, *, exc_context=None)\n\ +throw(type[,value[,tb]], /, *, exc_context=None)\n\ \n\ Raise exception in generator, return next yielded value or raise\n\ StopIteration.\n\ @@ -622,7 +622,8 @@ and may be removed in a future version of Python."); static PyObject * _gen_throw(PyGenObject *gen, int close_on_genexit, - PyObject *typ, PyObject *val, PyObject *tb) + PyObject *typ, PyObject *val, PyObject *tb, + PyObject *exc_context) { int8_t frame_state = FT_ATOMIC_LOAD_INT8_RELAXED(gen->gi_frame_state); do { @@ -680,7 +681,7 @@ _gen_throw(PyGenObject *gen, int close_on_genexit, /* Close the generator that we are currently iterating with 'yield from' or awaiting on with 'await'. */ ret = _gen_throw((PyGenObject *)yf, close_on_genexit, - typ, val, tb); + typ, val, tb, exc_context); tstate->current_frame = prev; frame->previous = NULL; } @@ -715,21 +716,53 @@ _gen_throw(PyGenObject *gen, int close_on_genexit, throw_here: assert(FT_ATOMIC_LOAD_INT8_RELAXED(gen->gi_frame_state) == FRAME_EXECUTING); + + /* + If an exception context is provided, set it (and restore it after) + + This ensures that after the generator handles the thrown-in exception, + _PyErr_GetTopmostException will return the provided exception context + instead of whatever exception was next in the stack. + */ + PyThreadState *tstate = NULL; + PyObject *prev_exc = NULL; + if (exc_context != NULL) { + tstate = _PyThreadState_GET(); + assert(tstate != NULL); + prev_exc = tstate->exc_info->exc_value; + tstate->exc_info->exc_value = Py_XNewRef(exc_context); + } if (gen_set_exception(typ, val, tb) < 0) { + if (tstate != NULL) { + Py_XSETREF(tstate->exc_info->exc_value, prev_exc); + } FT_ATOMIC_STORE_INT8_RELEASE(gen->gi_frame_state, frame_state); return NULL; } - return gen_throw_current_exception(gen); + + PyObject *result = gen_throw_current_exception(gen); + if (tstate != NULL) { + Py_XSETREF(tstate->exc_info->exc_value, prev_exc); + } + return result; } +/* +throw(...) method of builtins.generator instance + throw(value, /, *, exc_context=None) + throw(type[,value[,tb]], /, *, exc_context=None) + Raise exception in generator, return next yielded value or raise + StopIteration. +*/ static PyObject * -gen_throw(PyObject *op, PyObject *const *args, Py_ssize_t nargs) +gen_throw(PyObject *op, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) { PyGenObject *gen = _PyGen_CAST(op); PyObject *typ; PyObject *tb = NULL; PyObject *val = NULL; + PyObject *exc_context = NULL; if (!_PyArg_CheckPositional("throw", nargs, 1, 3)) { return NULL; @@ -742,6 +775,18 @@ gen_throw(PyObject *op, PyObject *const *args, Py_ssize_t nargs) return NULL; } } + + static const char * const _keywords[] = {"", "", "", "exc_context", NULL}; + static _PyArg_Parser _parser = { + .keywords = _keywords, + .fname = "throw", + }; + PyObject *argsbuf[4]; + args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 1, 3, 0, 1, argsbuf); + if (!args) { + return NULL; + } + typ = args[0]; if (nargs == 3) { val = args[1]; @@ -750,7 +795,19 @@ gen_throw(PyObject *op, PyObject *const *args, Py_ssize_t nargs) else if (nargs == 2) { val = args[1]; } - return _gen_throw(gen, 1, typ, val, tb); + + if (kwnames && PyTuple_GET_SIZE(kwnames)){ + exc_context = args[3]; + if (!Py_IsNone(exc_context) && !PyExceptionInstance_Check(exc_context)){ + PyErr_SetString(PyExc_TypeError, "exc_context must be an Exception object or None"); + return NULL; + } + } + /* default to current exception */ + if (exc_context == NULL){ + exc_context = _PyErr_GetTopmostException(_PyThreadState_GET())->exc_value; + } + return _gen_throw(gen, 1, typ, val, tb, exc_context); } @@ -1020,7 +1077,7 @@ PyDoc_STRVAR(sizeof__doc__, static PyMethodDef gen_methods[] = { {"send", gen_send, METH_O, send_doc}, - {"throw", _PyCFunction_CAST(gen_throw), METH_FASTCALL, throw_doc}, + {"throw", _PyCFunction_CAST(gen_throw), METH_FASTCALL|METH_KEYWORDS, throw_doc}, {"close", gen_close, METH_NOARGS, close_doc}, {"__sizeof__", gen_sizeof, METH_NOARGS, sizeof__doc__}, {"__class_getitem__", Py_GenericAlias, METH_O|METH_CLASS, PyDoc_STR("See PEP 585")}, @@ -1357,8 +1414,8 @@ PyDoc_STRVAR(coro_send_doc, return next iterated value or raise StopIteration."); PyDoc_STRVAR(coro_throw_doc, -"throw(value)\n\ -throw(type[,value[,traceback]])\n\ +"throw(value, /, *, exc_context=None)\n\ +throw(type[,value[,tb]], *, exc_context=None)\n\ \n\ Raise exception in coroutine, return next iterated value or raise\n\ StopIteration.\n\ @@ -1371,7 +1428,7 @@ PyDoc_STRVAR(coro_close_doc, static PyMethodDef coro_methods[] = { {"send", gen_send, METH_O, coro_send_doc}, - {"throw",_PyCFunction_CAST(gen_throw), METH_FASTCALL, coro_throw_doc}, + {"throw", _PyCFunction_CAST(gen_throw), METH_FASTCALL|METH_KEYWORDS, coro_throw_doc}, {"close", gen_close, METH_NOARGS, coro_close_doc}, {"__sizeof__", gen_sizeof, METH_NOARGS, sizeof__doc__}, {"__class_getitem__", Py_GenericAlias, METH_O|METH_CLASS, PyDoc_STR("See PEP 585")}, @@ -1461,10 +1518,10 @@ coro_wrapper_send(PyObject *self, PyObject *arg) } static PyObject * -coro_wrapper_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs) +coro_wrapper_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) { PyCoroWrapper *cw = _PyCoroWrapper_CAST(self); - return gen_throw((PyObject*)cw->cw_coroutine, args, nargs); + return gen_throw((PyObject*)cw->cw_coroutine, args, nargs, kwnames); } static PyObject * @@ -1484,8 +1541,8 @@ coro_wrapper_traverse(PyObject *self, visitproc visit, void *arg) static PyMethodDef coro_wrapper_methods[] = { {"send", coro_wrapper_send, METH_O, coro_send_doc}, - {"throw", _PyCFunction_CAST(coro_wrapper_throw), METH_FASTCALL, - coro_throw_doc}, + {"throw", _PyCFunction_CAST(coro_wrapper_throw), + METH_FASTCALL|METH_KEYWORDS, coro_throw_doc}, {"close", coro_wrapper_close, METH_NOARGS, coro_close_doc}, {NULL, NULL} /* Sentinel */ }; @@ -2024,7 +2081,7 @@ async_gen_asend_iternext(PyObject *ags) static PyObject * -async_gen_asend_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs) +async_gen_asend_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) { PyAsyncGenASend *o = _PyAsyncGenASend_CAST(self); @@ -2048,7 +2105,7 @@ async_gen_asend_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs) o->ags_gen->ag_running_async = 1; } - PyObject *result = gen_throw((PyObject*)o->ags_gen, args, nargs); + PyObject *result = gen_throw((PyObject*)o->ags_gen, args, nargs, kwnames); result = async_gen_unwrap_value(o->ags_gen, result); if (result == NULL) { @@ -2068,7 +2125,7 @@ async_gen_asend_close(PyObject *self, PyObject *args) Py_RETURN_NONE; } - PyObject *result = async_gen_asend_throw(self, &PyExc_GeneratorExit, 1); + PyObject *result = async_gen_asend_throw(self, &PyExc_GeneratorExit, 1, NULL); if (result == NULL) { if (PyErr_ExceptionMatches(PyExc_StopIteration) || PyErr_ExceptionMatches(PyExc_StopAsyncIteration) || @@ -2096,7 +2153,8 @@ async_gen_asend_finalize(PyObject *self) static PyMethodDef async_gen_asend_methods[] = { {"send", async_gen_asend_send, METH_O, send_doc}, - {"throw", _PyCFunction_CAST(async_gen_asend_throw), METH_FASTCALL, throw_doc}, + {"throw", _PyCFunction_CAST(async_gen_asend_throw), + METH_FASTCALL|METH_KEYWORDS, throw_doc}, {"close", async_gen_asend_close, METH_NOARGS, close_doc}, {NULL, NULL} /* Sentinel */ }; @@ -2349,7 +2407,7 @@ async_gen_athrow_send(PyObject *self, PyObject *arg) retval = _gen_throw((PyGenObject *)gen, 0, /* Do not close generator when PyExc_GeneratorExit is passed */ - PyExc_GeneratorExit, NULL, NULL); + PyExc_GeneratorExit, NULL, NULL, NULL); if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) { Py_DECREF(retval); @@ -2359,7 +2417,7 @@ async_gen_athrow_send(PyObject *self, PyObject *arg) retval = _gen_throw((PyGenObject *)gen, 0, /* Do not close generator when PyExc_GeneratorExit is passed */ - o->agt_typ, o->agt_val, o->agt_tb); + o->agt_typ, o->agt_val, o->agt_tb, NULL); retval = async_gen_unwrap_value(o->agt_gen, retval); } if (retval == NULL) { @@ -2417,7 +2475,7 @@ async_gen_athrow_send(PyObject *self, PyObject *arg) static PyObject * -async_gen_athrow_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs) +async_gen_athrow_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) { PyAsyncGenAThrow *o = _PyAsyncGenAThrow_CAST(self); @@ -2448,7 +2506,7 @@ async_gen_athrow_throw(PyObject *self, PyObject *const *args, Py_ssize_t nargs) o->agt_gen->ag_running_async = 1; } - PyObject *retval = gen_throw((PyObject*)o->agt_gen, args, nargs); + PyObject *retval = gen_throw((PyObject*)o->agt_gen, args, nargs, kwnames); if (o->agt_typ) { retval = async_gen_unwrap_value(o->agt_gen, retval); if (retval == NULL) { @@ -2501,7 +2559,7 @@ async_gen_athrow_close(PyObject *self, PyObject *args) Py_RETURN_NONE; } PyObject *result = async_gen_athrow_throw((PyObject*)agt, - &PyExc_GeneratorExit, 1); + &PyExc_GeneratorExit, 1, NULL); if (result == NULL) { if (PyErr_ExceptionMatches(PyExc_StopIteration) || PyErr_ExceptionMatches(PyExc_StopAsyncIteration) || @@ -2532,7 +2590,7 @@ async_gen_athrow_finalize(PyObject *op) static PyMethodDef async_gen_athrow_methods[] = { {"send", async_gen_athrow_send, METH_O, send_doc}, {"throw", _PyCFunction_CAST(async_gen_athrow_throw), - METH_FASTCALL, throw_doc}, + METH_FASTCALL|METH_KEYWORDS, throw_doc}, {"close", async_gen_athrow_close, METH_NOARGS, close_doc}, {NULL, NULL} /* Sentinel */ }; diff --git a/Python/errors.c b/Python/errors.c index 48b03e5fd714b18..a1014bc6f1e0d41 100644 --- a/Python/errors.c +++ b/Python/errors.c @@ -121,7 +121,6 @@ _PyErr_GetTopmostException(PyThreadState *tstate) { exc_info = exc_info->previous_item; } - assert(!Py_IsNone(exc_info->exc_value)); return exc_info; } @@ -607,7 +606,7 @@ PyErr_GetHandledException(void) void _PyErr_SetHandledException(PyThreadState *tstate, PyObject *exc) { - Py_XSETREF(tstate->exc_info->exc_value, Py_XNewRef(exc == Py_None ? NULL : exc)); + Py_XSETREF(tstate->exc_info->exc_value, Py_XNewRef(exc ? exc : Py_None)); } void