bpo-41756: Refactor gen_send_ex(). (GH-22330)
diff --git a/Objects/genobject.c b/Objects/genobject.c
index 24aca98..f0943ae 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -136,13 +136,15 @@
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 @@
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 @@
}
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 @@
"just-started async generator";
}
PyErr_SetString(PyExc_TypeError, msg);
- return NULL;
+ return PYGEN_ERROR;
}
} else {
/* Push arg onto the frame's value stack */
@@ -229,69 +224,77 @@
/* 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);
+ if (result) {
+ if (!_PyFrameHasCompleted(f)) {
+ *presult = result;
+ return PYGEN_NEXT;
+ }
+ 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 (PyErr_ExceptionMatches(PyExc_StopIteration)) {
+ const char *msg = "generator raised StopIteration";
+ if (PyCoro_CheckExact(gen)) {
+ msg = "coroutine raised StopIteration";
}
- 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 if (PyAsyncGen_CheckExact(gen)) {
+ msg = "async generator raised StopIteration";
}
- else {
- Py_CLEAR(result);
- }
+ _PyErr_FormatFromCause(PyExc_RuntimeError, "%s", msg);
+ }
+ 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);
+ }
+ }
+
+ /* 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 {
- /* 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);
- }
+ _PyGen_SetStopIterationValue(result);
}
+ 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 (PyAsyncGen_CheckExact(gen)) {
- msg = "async generator raised StopIteration";
- }
- _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);
- }
-
return result;
}
@@ -299,22 +302,16 @@
"send(arg) -> send 'arg' into generator,\n\
return next yielded value or raise StopIteration.");
+static PyObject *
+gen_send(PyGenObject *gen, PyObject *arg)
+{
+ return gen_send_ex(gen, arg, 0, 0);
+}
+
PyObject *
_PyGen_Send(PyGenObject *gen, PyObject *arg)
{
- return gen_send_ex(gen, arg, 0, 0, NULL);
-}
-
-PySendResult
-PyGen_Send(PyGenObject *gen, PyObject *arg, PyObject **result)
-{
- 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 @@
}
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->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 @@
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 @@
}
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 @@
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 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 @@
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 @@
}
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 @@
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 {