bpo-31861: Add aiter and anext to builtins (#23847)

Co-authored-by: jab <jab@users.noreply.github.com>
Co-authored-by: Daniel Pope <mauve@mauveweb.co.uk>
Co-authored-by: Justin Wang <justin39@gmail.com>

1 files changed, 31 insertions, 0 deletions

diff --git a/Objects/abstract.c b/Objects/abstract.c
index 4cd59100d..fcfe2db 100644
--- a/Objects/abstract.c
+++ b/Objects/abstract.c
@@ -2738,6 +2738,26 @@ PyObject_GetIter(PyObject *o)
     }
 }
 
+PyObject *
+PyObject_GetAiter(PyObject *o) {
+    PyTypeObject *t = Py_TYPE(o);
+    unaryfunc f;
+
+    if (t->tp_as_async == NULL || t->tp_as_async->am_aiter == NULL) {
+        return type_error("'%.200s' object is not an AsyncIterable", o);
+    }
+    f = t->tp_as_async->am_aiter;
+    PyObject *it = (*f)(o);
+    if (it != NULL && !PyAiter_Check(it)) {
+        PyErr_Format(PyExc_TypeError,
+                     "aiter() returned non-AsyncIterator of type '%.100s'",
+                     Py_TYPE(it)->tp_name);
+        Py_DECREF(it);
+        it = NULL;
+    }
+    return it;
+}
+
 int
 PyIter_Check(PyObject *obj)
 {
@@ -2746,6 +2766,17 @@ PyIter_Check(PyObject *obj)
             tp->tp_iternext != &_PyObject_NextNotImplemented);
 }
 
+int
+PyAiter_Check(PyObject *obj)
+{
+    PyTypeObject *tp = Py_TYPE(obj);
+    return (tp->tp_as_async != NULL &&
+            tp->tp_as_async->am_aiter != NULL &&
+            tp->tp_as_async->am_aiter != &_PyObject_NextNotImplemented &&
+            tp->tp_as_async->am_anext != NULL &&
+            tp->tp_as_async->am_anext != &_PyObject_NextNotImplemented);
+}
+
 /* Return next item.
  * If an error occurs, return NULL.  PyErr_Occurred() will be true.
  * If the iteration terminates normally, return NULL and clear the