bpo-44963: Implement send() and throw() methods for anext_awaitable objects (GH-27955)
Co-authored-by: Yury Selivanov <yury@edgedb.com>
(cherry picked from commit 533e725821b15e2df2cd4479a34597c1d8faf616)
Co-authored-by: Pablo Galindo Salgado <Pablogsal@gmail.com>
diff --git a/Objects/iterobject.c b/Objects/iterobject.c
index 6961fc3..e493e41 100644
--- a/Objects/iterobject.c
+++ b/Objects/iterobject.c
@@ -314,6 +314,36 @@ anextawaitable_traverse(anextawaitableobject *obj, visitproc visit, void *arg)
}
static PyObject *
+anextawaitable_getiter(anextawaitableobject *obj)
+{
+ assert(obj->wrapped != NULL);
+ PyObject *awaitable = _PyCoro_GetAwaitableIter(obj->wrapped);
+ if (awaitable == NULL) {
+ return NULL;
+ }
+ if (Py_TYPE(awaitable)->tp_iternext == NULL) {
+ /* _PyCoro_GetAwaitableIter returns a Coroutine, a Generator,
+ * or an iterator. Of these, only coroutines lack tp_iternext.
+ */
+ assert(PyCoro_CheckExact(awaitable));
+ unaryfunc getter = Py_TYPE(awaitable)->tp_as_async->am_await;
+ PyObject *new_awaitable = getter(awaitable);
+ if (new_awaitable == NULL) {
+ Py_DECREF(awaitable);
+ return NULL;
+ }
+ Py_SETREF(awaitable, new_awaitable);
+ if (!PyIter_Check(awaitable)) {
+ PyErr_SetString(PyExc_TypeError,
+ "__await__ returned a non-iterable");
+ Py_DECREF(awaitable);
+ return NULL;
+ }
+ }
+ return awaitable;
+}
+
+static PyObject *
anextawaitable_iternext(anextawaitableobject *obj)
{
/* Consider the following class:
@@ -336,30 +366,10 @@ anextawaitable_iternext(anextawaitableobject *obj)
* Then `await anext(gen)` can just call
* gen.__anext__().__next__()
*/
- assert(obj->wrapped != NULL);
- PyObject *awaitable = _PyCoro_GetAwaitableIter(obj->wrapped);
+ PyObject *awaitable = anextawaitable_getiter(obj);
if (awaitable == NULL) {
return NULL;
}
- if (Py_TYPE(awaitable)->tp_iternext == NULL) {
- /* _PyCoro_GetAwaitableIter returns a Coroutine, a Generator,
- * or an iterator. Of these, only coroutines lack tp_iternext.
- */
- assert(PyCoro_CheckExact(awaitable));
- unaryfunc getter = Py_TYPE(awaitable)->tp_as_async->am_await;
- PyObject *new_awaitable = getter(awaitable);
- if (new_awaitable == NULL) {
- Py_DECREF(awaitable);
- return NULL;
- }
- Py_SETREF(awaitable, new_awaitable);
- if (Py_TYPE(awaitable)->tp_iternext == NULL) {
- PyErr_SetString(PyExc_TypeError,
- "__await__ returned a non-iterable");
- Py_DECREF(awaitable);
- return NULL;
- }
- }
PyObject *result = (*Py_TYPE(awaitable)->tp_iternext)(awaitable);
Py_DECREF(awaitable);
if (result != NULL) {
@@ -371,6 +381,70 @@ anextawaitable_iternext(anextawaitableobject *obj)
return NULL;
}
+
+static PyObject *
+anextawaitable_proxy(anextawaitableobject *obj, char *meth, PyObject *arg) {
+ PyObject *awaitable = anextawaitable_getiter(obj);
+ if (awaitable == NULL) {
+ return NULL;
+ }
+ PyObject *ret = PyObject_CallMethod(awaitable, meth, "O", arg);
+ Py_DECREF(awaitable);
+ if (ret != NULL) {
+ return ret;
+ }
+ if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
+ /* `anextawaitableobject` is only used by `anext()` when
+ * a default value is provided. So when we have a StopAsyncIteration
+ * exception we replace it with a `StopIteration(default)`, as if
+ * it was the return value of `__anext__()` coroutine.
+ */
+ _PyGen_SetStopIterationValue(obj->default_value);
+ }
+ return NULL;
+}
+
+
+static PyObject *
+anextawaitable_send(anextawaitableobject *obj, PyObject *arg) {
+ return anextawaitable_proxy(obj, "send", arg);
+}
+
+
+static PyObject *
+anextawaitable_throw(anextawaitableobject *obj, PyObject *arg) {
+ return anextawaitable_proxy(obj, "throw", arg);
+}
+
+
+static PyObject *
+anextawaitable_close(anextawaitableobject *obj, PyObject *arg) {
+ return anextawaitable_proxy(obj, "close", arg);
+}
+
+
+PyDoc_STRVAR(send_doc,
+"send(arg) -> send 'arg' into the wrapped iterator,\n\
+return next yielded value or raise StopIteration.");
+
+
+PyDoc_STRVAR(throw_doc,
+"throw(typ[,val[,tb]]) -> raise exception in the wrapped iterator,\n\
+return next yielded value or raise StopIteration.");
+
+
+PyDoc_STRVAR(close_doc,
+"close() -> raise GeneratorExit inside generator.");
+
+
+static PyMethodDef anextawaitable_methods[] = {
+ {"send",(PyCFunction)anextawaitable_send, METH_O, send_doc},
+ {"throw",(PyCFunction)anextawaitable_throw, METH_VARARGS, throw_doc},
+ {"close",(PyCFunction)anextawaitable_close, METH_VARARGS, close_doc},
+ {NULL, NULL} /* Sentinel */
+};
+
+
static PyAsyncMethods anextawaitable_as_async = {
PyObject_SelfIter, /* am_await */
0, /* am_aiter */
@@ -407,7 +481,7 @@ PyTypeObject _PyAnextAwaitable_Type = {
0, /* tp_weaklistoffset */
PyObject_SelfIter, /* tp_iter */
(unaryfunc)anextawaitable_iternext, /* tp_iternext */
- 0, /* tp_methods */
+ anextawaitable_methods, /* tp_methods */
};
PyObject *