Issue 24342: Let wrapper set by sys.set_coroutine_wrapper fail gracefully

1 files changed, 30 insertions, 9 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index bb2c0b96a5..2a1db17b45 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -3921,7 +3921,6 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals,
 
     if (co->co_flags & CO_GENERATOR) {
         PyObject *gen;
-        PyObject *coroutine_wrapper;
 
         /* Don't need to keep the reference to f_back, it will be set
          * when the generator is resumed. */
@@ -3935,14 +3934,9 @@ _PyEval_EvalCodeWithName(PyObject *_co, PyObject *globals, PyObject *locals,
         if (gen == NULL)
             return NULL;
 
-        if (co->co_flags & (CO_COROUTINE | CO_ITERABLE_COROUTINE)) {
-            coroutine_wrapper = _PyEval_GetCoroutineWrapper();
-            if (coroutine_wrapper != NULL) {
-                PyObject *wrapped =
-                            PyObject_CallFunction(coroutine_wrapper, "N", gen);
-                gen = wrapped;
-            }
-        }
+        if (co->co_flags & (CO_COROUTINE | CO_ITERABLE_COROUTINE))
+            return _PyEval_ApplyCoroutineWrapper(gen);
+
         return gen;
     }
 
@@ -4408,6 +4402,33 @@ _PyEval_GetCoroutineWrapper(void)
 }
 
 PyObject *
+_PyEval_ApplyCoroutineWrapper(PyObject *gen)
+{
+    PyObject *wrapped;
+    PyThreadState *tstate = PyThreadState_GET();
+    PyObject *wrapper = tstate->coroutine_wrapper;
+
+    if (tstate->in_coroutine_wrapper) {
+        assert(wrapper != NULL);
+        PyErr_Format(PyExc_RuntimeError,
+                     "coroutine wrapper %.150R attempted "
+                     "to recursively wrap %.150R",
+                     wrapper,
+                     gen);
+        return NULL;
+    }
+
+    if (wrapper == NULL) {
+        return gen;
+    }
+
+    tstate->in_coroutine_wrapper = 1;
+    wrapped = PyObject_CallFunction(wrapper, "N", gen);
+    tstate->in_coroutine_wrapper = 0;
+    return wrapped;
+}
+
+PyObject *
 PyEval_GetBuiltins(void)
 {
     PyFrameObject *current_frame = PyEval_GetFrame();