Add exc_context parameter to {gen/coro}.throw()

This solves the exception context leaking into the generator by allowing
an alternate exception context to be passed in. It also removes the need
for an exposed function that directly manipulates the exception state
since the details are handled in the C implementation.

Author: pR0Ps

Lib/contextlib.py

@@ -235,11 +235,10 @@ def __exit__(self, typ, value, traceback):
                 # 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.
+                # 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)
+                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

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 */
 };