bpo-41756: Refactor gen_send_ex(). (GH-22330)

Author: serhiy-storchaka

Objects/genobject.c

@@ -136,13 +136,15 @@ gen_dealloc(PyGenObject *gen)
     PyObject_GC_Del(gen);
 }
 
-static PyObject *
-gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing, int *is_return_value)
+static PySendResult
+gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
+             int exc, int closing)
 {
     PyThreadState *tstate = _PyThreadState_GET();
     PyFrameObject *f = gen->gi_frame;
     PyObject *result;
 
+    *presult = NULL;
     if (f != NULL && _PyFrame_IsExecuting(f)) {
         const char *msg = "generator already executing";
         if (PyCoro_CheckExact(gen)) {
@@ -152,7 +154,7 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing, int *is_retur
             msg = "async generator already executing";
         }
         PyErr_SetString(PyExc_ValueError, msg);
-        return NULL;
+        return PYGEN_ERROR;
     }
     if (f == NULL || _PyFrameHasCompleted(f)) {
         if (PyCoro_CheckExact(gen) && !closing) {
@@ -165,19 +167,12 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing, int *is_retur
         }
         else if (arg && !exc) {
             /* `gen` is an exhausted generator:
-               only set exception if called from send(). */
-            if (PyAsyncGen_CheckExact(gen)) {
-                PyErr_SetNone(PyExc_StopAsyncIteration);
-            }
-            else {
-                if (is_return_value != NULL) {
-                    *is_return_value = 1;
-                    Py_RETURN_NONE;
-                }
-                PyErr_SetNone(PyExc_StopIteration);
-            }
+               only return value if called from send(). */
+            *presult = Py_None;
+            Py_INCREF(*presult);
+            return PYGEN_RETURN;
         }
-        return NULL;
+        return PYGEN_ERROR;
     }
 
     assert(_PyFrame_IsRunnable(f));
@@ -193,7 +188,7 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing, int *is_retur
                       "just-started async generator";
             }
             PyErr_SetString(PyExc_TypeError, msg);
-            return NULL;
+            return PYGEN_ERROR;
         }
     } else {
         /* Push arg onto the frame's value stack */
@@ -229,92 +224,94 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing, int *is_retur
 
     /* If the generator just returned (as opposed to yielding), signal
      * that the generator is exhausted. */
-    if (result &&  _PyFrameHasCompleted(f)) {
-        if (result == Py_None) {
-            /* Delay exception instantiation if we can */
-            if (PyAsyncGen_CheckExact(gen)) {
-                PyErr_SetNone(PyExc_StopAsyncIteration);
-                Py_CLEAR(result);
-            }
-            else if (arg) {
-                if (is_return_value != NULL) {
-                    *is_return_value = 1;
-                }
-                else {
-                    /* Set exception if not called by gen_iternext() */
-                    PyErr_SetNone(PyExc_StopIteration);
-                    Py_CLEAR(result);
-                }
-            }
-            else {
-                Py_CLEAR(result);
-            }
+    if (result) {
+        if (!_PyFrameHasCompleted(f)) {
+            *presult = result;
+            return PYGEN_NEXT;
         }
-        else {
-            /* Async generators cannot return anything but None */
-            assert(!PyAsyncGen_CheckExact(gen));
-            if (is_return_value != NULL) {
-                *is_return_value = 1;
-            }
-            else {
-                _PyGen_SetStopIterationValue(result);
-                Py_CLEAR(result);
-            }
+        assert(result == Py_None || !PyAsyncGen_CheckExact(gen));
+        if (result == Py_None && !PyAsyncGen_CheckExact(gen) && !arg) {
+            /* Return NULL if called by gen_iternext() */
+            Py_CLEAR(result);
         }
     }
-    else if (!result && PyErr_ExceptionMatches(PyExc_StopIteration)) {
-        const char *msg = "generator raised StopIteration";
-        if (PyCoro_CheckExact(gen)) {
-            msg = "coroutine raised StopIteration";
+    else {
+        if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
+            const char *msg = "generator raised StopIteration";
+            if (PyCoro_CheckExact(gen)) {
+                msg = "coroutine raised StopIteration";
+            }
+            else if (PyAsyncGen_CheckExact(gen)) {
+                msg = "async generator raised StopIteration";
+            }
+            _PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
         }
-        else if (PyAsyncGen_CheckExact(gen)) {
-            msg = "async generator raised StopIteration";
+        else if (PyAsyncGen_CheckExact(gen) &&
+                PyErr_ExceptionMatches(PyExc_StopAsyncIteration))
+        {
+            /* code in `gen` raised a StopAsyncIteration error:
+               raise a RuntimeError.
+            */
+            const char *msg = "async generator raised StopAsyncIteration";
+            _PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
         }
-        _PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
-
-    }
-    else if (!result && PyAsyncGen_CheckExact(gen) &&
-             PyErr_ExceptionMatches(PyExc_StopAsyncIteration))
-    {
-        /* code in `gen` raised a StopAsyncIteration error:
-           raise a RuntimeError.
-        */
-        const char *msg = "async generator raised StopAsyncIteration";
-        _PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
     }
 
-    if ((is_return_value && *is_return_value) || !result ||  _PyFrameHasCompleted(f)) {
-        /* generator can't be rerun, so release the frame */
-        /* first clean reference cycle through stored exception traceback */
-        _PyErr_ClearExcState(&gen->gi_exc_state);
-        gen->gi_frame->f_gen = NULL;
-        gen->gi_frame = NULL;
-        Py_DECREF(f);
-    }
+    /* generator can't be rerun, so release the frame */
+    /* first clean reference cycle through stored exception traceback */
+    _PyErr_ClearExcState(&gen->gi_exc_state);
+    gen->gi_frame->f_gen = NULL;
+    gen->gi_frame = NULL;
+    Py_DECREF(f);
+
+    *presult = result;
+    return result ? PYGEN_RETURN : PYGEN_ERROR;
+}
+
+PySendResult
+PyGen_Send(PyGenObject *gen, PyObject *arg, PyObject **result)
+{
+    assert(PyGen_CheckExact(gen) || PyCoro_CheckExact(gen));
+    assert(result != NULL);
+    assert(arg != NULL);
+
+    return gen_send_ex2(gen, arg, result, 0, 0);
+}
 
+static PyObject *
+gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
+{
+    PyObject *result;
+    if (gen_send_ex2(gen, arg, &result, exc, closing) == PYGEN_RETURN) {
+        if (PyAsyncGen_CheckExact(gen)) {
+            assert(result == Py_None);
+            PyErr_SetNone(PyExc_StopAsyncIteration);
+        }
+        else if (result == Py_None) {
+            PyErr_SetNone(PyExc_StopIteration);
+        }
+        else {
+            _PyGen_SetStopIterationValue(result);
+        }
+        Py_CLEAR(result);
+    }
     return result;
 }
 
 PyDoc_STRVAR(send_doc,
 "send(arg) -> send 'arg' into generator,\n\
 return next yielded value or raise StopIteration.");
 
-PyObject *
-_PyGen_Send(PyGenObject *gen, PyObject *arg)
+static PyObject *
+gen_send(PyGenObject *gen, PyObject *arg)
 {
-    return gen_send_ex(gen, arg, 0, 0, NULL);
+    return gen_send_ex(gen, arg, 0, 0);
 }
 
-PySendResult
-PyGen_Send(PyGenObject *gen, PyObject *arg, PyObject **result)
+PyObject *
+_PyGen_Send(PyGenObject *gen, PyObject *arg)
 {
-    assert(result != NULL);
-
-    int is_return_value = 0;
-    if ((*result = gen_send_ex(gen, arg, 0, 0, &is_return_value)) == NULL) {
-        return PYGEN_ERROR;
-    }
-    return is_return_value ? PYGEN_RETURN : PYGEN_NEXT;
+    return gen_send(gen, arg);
 }
 
 PyDoc_STRVAR(close_doc,
@@ -396,7 +393,7 @@ gen_close(PyGenObject *gen, PyObject *args)
     }
     if (err == 0)
         PyErr_SetNone(PyExc_GeneratorExit);
-    retval = gen_send_ex(gen, Py_None, 1, 1, NULL);
+    retval = gen_send_ex(gen, Py_None, 1, 1);
     if (retval) {
         const char *msg = "generator ignored GeneratorExit";
         if (PyCoro_CheckExact(gen)) {
@@ -444,7 +441,7 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
             gen->gi_frame->f_state = state;
             Py_DECREF(yf);
             if (err < 0)
-                return gen_send_ex(gen, Py_None, 1, 0, NULL);
+                return gen_send_ex(gen, Py_None, 1, 0);
             goto throw_here;
         }
         if (PyGen_CheckExact(yf) || PyCoro_CheckExact(yf)) {
@@ -496,10 +493,10 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
             assert(gen->gi_frame->f_lasti >= 0);
             gen->gi_frame->f_lasti += sizeof(_Py_CODEUNIT);
             if (_PyGen_FetchStopIterationValue(&val) == 0) {
-                ret = gen_send_ex(gen, val, 0, 0, NULL);
+                ret = gen_send(gen, val);
                 Py_DECREF(val);
             } else {
-                ret = gen_send_ex(gen, Py_None, 1, 0, NULL);
+                ret = gen_send_ex(gen, Py_None, 1, 0);
             }
         }
         return ret;
@@ -553,7 +550,7 @@ _gen_throw(PyGenObject *gen, int close_on_genexit,
     }
 
     PyErr_Restore(typ, val, tb);
-    return gen_send_ex(gen, Py_None, 1, 0, NULL);
+    return gen_send_ex(gen, Py_None, 1, 0);
 
 failed_throw:
     /* Didn't use our arguments, so restore their original refcounts */
@@ -582,7 +579,15 @@ gen_throw(PyGenObject *gen, PyObject *args)
 static PyObject *
 gen_iternext(PyGenObject *gen)
 {
-    return gen_send_ex(gen, NULL, 0, 0, NULL);
+    PyObject *result;
+    assert(PyGen_CheckExact(gen) || PyCoro_CheckExact(gen));
+    if (gen_send_ex2(gen, NULL, &result, 0, 0) == PYGEN_RETURN) {
+        if (result != Py_None) {
+            _PyGen_SetStopIterationValue(result);
+        }
+        Py_CLEAR(result);
+    }
+    return result;
 }
 
 /*
@@ -767,7 +772,7 @@ static PyMemberDef gen_memberlist[] = {
 };
 
 static PyMethodDef gen_methods[] = {
-    {"send",(PyCFunction)_PyGen_Send, METH_O, send_doc},
+    {"send",(PyCFunction)gen_send, METH_O, send_doc},
     {"throw",(PyCFunction)gen_throw, METH_VARARGS, throw_doc},
     {"close",(PyCFunction)gen_close, METH_NOARGS, close_doc},
     {NULL, NULL}        /* Sentinel */
@@ -1082,13 +1087,13 @@ coro_wrapper_dealloc(PyCoroWrapper *cw)
 static PyObject *
 coro_wrapper_iternext(PyCoroWrapper *cw)
 {
-    return gen_send_ex((PyGenObject *)cw->cw_coroutine, NULL, 0, 0, NULL);
+    return gen_iternext((PyGenObject *)cw->cw_coroutine);
 }
 
 static PyObject *
 coro_wrapper_send(PyCoroWrapper *cw, PyObject *arg)
 {
-    return gen_send_ex((PyGenObject *)cw->cw_coroutine, arg, 0, 0, NULL);
+    return gen_send((PyGenObject *)cw->cw_coroutine, arg);
 }
 
 static PyObject *
@@ -1601,7 +1606,7 @@ async_gen_asend_send(PyAsyncGenASend *o, PyObject *arg)
     }
 
     o->ags_gen->ag_running_async = 1;
-    result = gen_send_ex((PyGenObject*)o->ags_gen, arg, 0, 0, NULL);
+    result = gen_send((PyGenObject*)o->ags_gen, arg);
     result = async_gen_unwrap_value(o->ags_gen, result);
 
     if (result == NULL) {
@@ -1957,7 +1962,7 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
 
     assert(o->agt_state == AWAITABLE_STATE_ITER);
 
-    retval = gen_send_ex((PyGenObject *)gen, arg, 0, 0, NULL);
+    retval = gen_send((PyGenObject *)gen, arg);
     if (o->agt_args) {
         return async_gen_unwrap_value(o->agt_gen, retval);
     } else {