bpo-44963: Implement send() and throw() methods for anext_awaitable objects (GH-27955)

Co-authored-by: Yury Selivanov <yury@edgedb.com>

1 files changed, 96 insertions, 22 deletions

diff --git a/Objects/iterobject.c b/Objects/iterobject.c
index 6961fc3b4a..e493e41131 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 *