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>
2021-03-23 18:47:21 -04:00 · 2021-03-23 18:47:21 -04:00 · f0a6fde882
parent 94faa0724f
commit f0a6fde882
10 changed files with 373 additions and 28 deletions
--- a/Doc/library/functions.rst
+++ b/Doc/library/functions.rst
@ -12,31 +12,31 @@ are always available.  They are listed here in alphabetical order.
 +=========================+=======================+=======================+=========================+
 | |  **A**                | |  **E**              | |  **L**              | |  **R**                |
 | |  :func:`abs`          | |  :func:`enumerate`  | |  :func:`len`        | |  |func-range|_        |
-| |  :func:`all`          | |  :func:`eval`       | |  |func-list|_       | |  :func:`repr`         |
-| |  :func:`any`          | |  :func:`exec`       | |  :func:`locals`     | |  :func:`reversed`     |
-| |  :func:`ascii`        | |                     | |                     | |  :func:`round`        |
-| |                       | |  **F**              | |  **M**              | |                       |
-| |  **B**                | |  :func:`filter`     | |  :func:`map`        | |  **S**                |
-| |  :func:`bin`          | |  :func:`float`      | |  :func:`max`        | |  |func-set|_          |
-| |  :func:`bool`         | |  :func:`format`     | |  |func-memoryview|_ | |  :func:`setattr`      |
-| |  :func:`breakpoint`   | |  |func-frozenset|_  | |  :func:`min`        | |  :func:`slice`        |
-| |  |func-bytearray|_    | |                     | |                     | |  :func:`sorted`       |
-| |  |func-bytes|_        | |  **G**              | |  **N**              | |  :func:`staticmethod` |
-| |                       | |  :func:`getattr`    | |  :func:`next`       | |  |func-str|_          |
-| |  **C**                | |  :func:`globals`    | |                     | |  :func:`sum`          |
-| |  :func:`callable`     | |                     | |  **O**              | |  :func:`super`        |
-| |  :func:`chr`          | |  **H**              | |  :func:`object`     | |                       |
-| |  :func:`classmethod`  | |  :func:`hasattr`    | |  :func:`oct`        | |  **T**                |
-| |  :func:`compile`      | |  :func:`hash`       | |  :func:`open`       | |  |func-tuple|_        |
-| |  :func:`complex`      | |  :func:`help`       | |  :func:`ord`        | |  :func:`type`         |
-| |                       | |  :func:`hex`        | |                     | |                       |
-| |  **D**                | |                     | |  **P**              | |  **V**                |
-| |  :func:`delattr`      | |  **I**              | |  :func:`pow`        | |  :func:`vars`         |
-| |  |func-dict|_         | |  :func:`id`         | |  :func:`print`      | |                       |
-| |  :func:`dir`          | |  :func:`input`      | |  :func:`property`   | |  **Z**                |
-| |  :func:`divmod`       | |  :func:`int`        | |                     | |  :func:`zip`          |
-| |                       | |  :func:`isinstance` | |                     | |                       |
-| |                       | |  :func:`issubclass` | |                     | |  **_**                |
+| |  :func:`aiter`        | |  :func:`eval`       | |  |func-list|_       | |  :func:`repr`         |
+| |  :func:`all`          | |  :func:`exec`       | |  :func:`locals`     | |  :func:`reversed`     |
+| |  :func:`any`          | |                     | |                     | |  :func:`round`        |
+| |  :func:`anext`        | |  **F**              | |  **M**              | |                       |
+| |  :func:`ascii`        | |  :func:`filter`     | |  :func:`map`        | |  **S**                |
+| |                       | |  :func:`float`      | |  :func:`max`        | |  |func-set|_          |
+| |  **B**                | |  :func:`format`     | |  |func-memoryview|_ | |  :func:`setattr`      |
+| |  :func:`bin`          | |  |func-frozenset|_  | |  :func:`min`        | |  :func:`slice`        |
+| |  :func:`bool`         | |                     | |                     | |  :func:`sorted`       |
+| |  :func:`breakpoint`   | |  **G**              | |  **N**              | |  :func:`staticmethod` |
+| |  |func-bytearray|_    | |  :func:`getattr`    | |  :func:`next`       | |  |func-str|_          |
+| |  |func-bytes|_        | |  :func:`globals`    | |                     | |  :func:`sum`          |
+| |                       | |                     | |  **O**              | |  :func:`super`        |
+| |  **C**                | |  **H**              | |  :func:`object`     | |                       |
+| |  :func:`callable`     | |  :func:`hasattr`    | |  :func:`oct`        | |  **T**                |
+| |  :func:`chr`          | |  :func:`hash`       | |  :func:`open`       | |  |func-tuple|_        |
+| |  :func:`classmethod`  | |  :func:`help`       | |  :func:`ord`        | |  :func:`type`         |
+| |  :func:`compile`      | |  :func:`hex`        | |                     | |                       |
+| |  :func:`complex`      | |                     | |  **P**              | |  **V**                |
+| |                       | |  **I**              | |  :func:`pow`        | |  :func:`vars`         |
+| |  **D**                | |  :func:`id`         | |  :func:`print`      | |                       |
+| |  :func:`delattr`      | |  :func:`input`      | |  :func:`property`   | |  **Z**                |
+| |  |func-dict|_         | |  :func:`int`        | |                     | |  :func:`zip`          |
+| |  :func:`dir`          | |  :func:`isinstance` | |                     | |                       |
+| |  :func:`divmod`       | |  :func:`issubclass` | |                     | |  **_**                |
 | |                       | |  :func:`iter`       | |                     | |  :func:`__import__`   |
 +-------------------------+-----------------------+-----------------------+-------------------------+

@ -61,6 +61,17 @@ are always available.  They are listed here in alphabetical order.
   If the argument is a complex number, its magnitude is returned.


+.. function:: aiter(async_iterable)
+
+   Return an :term:`asynchronous iterator` for an :term:`asynchronous iterable`.
+   Equivalent to calling ``x.__aiter__()``.
+
+   ``aiter(x)`` itself has an ``__aiter__()`` method that returns ``x``,
+   so ``aiter(aiter(x))`` is the same as ``aiter(x)``.
+
+   Note: Unlike :func:`iter`, :func:`aiter` has no 2-argument variant.
+
+
 .. function:: all(iterable)

   Return ``True`` if all elements of the *iterable* are true (or if the iterable
@ -73,6 +84,20 @@ are always available.  They are listed here in alphabetical order.
          return True


+.. awaitablefunction:: anext(async_iterator[, default])
+
+   When awaited, return the next item from the given :term:`asynchronous
+   iterator`, or *default* if given and the iterator is exhausted.
+
+   This is the async variant of the :func:`next` builtin, and behaves
+   similarly.
+
+   This calls the :meth:`~object.__anext__` method of *async_iterator*,
+   returning an :term:`awaitable`. Awaiting this returns the next value of the
+   iterator. If *default* is given, it is returned if the iterator is exhausted,
+   otherwise :exc:`StopAsyncIteration` is raised.
+
+
 .. function:: any(iterable)

   Return ``True`` if any element of the *iterable* is true.  If the iterable
--- a/Doc/whatsnew/3.10.rst
+++ b/Doc/whatsnew/3.10.rst
@ -588,6 +588,11 @@ Other Language Changes
  ``__globals__["__builtins__"]`` if it exists, else from the current builtins.
  (Contributed by Mark Shannon in :issue:`42990`.)

+* Two new builtin functions -- :func:`aiter` and :func:`anext` have been added
+  to provide asynchronous counterparts to :func:`iter` and :func:`next`,
+  respectively.
+  (Contributed by Joshua Bronson, Daniel Pope, and Justin Wang in :issue:`31861`.)
+

 New Modules
 ===========
--- a/Include/abstract.h
+++ b/Include/abstract.h
@ -324,11 +324,21 @@ PyAPI_FUNC(PyObject *) PyObject_Format(PyObject *obj,
   returns itself. */
 PyAPI_FUNC(PyObject *) PyObject_GetIter(PyObject *);

+/* Takes an AsyncIterable object and returns an AsyncIterator for it.
+   This is typically a new iterator but if the argument is an AsyncIterator,
+   this returns itself. */
+PyAPI_FUNC(PyObject *) PyObject_GetAiter(PyObject *);
+
 /* Returns non-zero if the object 'obj' provides iterator protocols, and 0 otherwise.

   This function always succeeds. */
 PyAPI_FUNC(int) PyIter_Check(PyObject *);

+/* Returns non-zero if the object 'obj' provides AsyncIterator protocols, and 0 otherwise.
+
+   This function always succeeds. */
+PyAPI_FUNC(int) PyAiter_Check(PyObject *);
+
 /* Takes an iterator object and calls its tp_iternext slot,
   returning the next value.

--- a/Lib/test/test_asyncgen.py
+++ b/Lib/test/test_asyncgen.py
@ -372,6 +372,88 @@ class AsyncGenAsyncioTest(unittest.TestCase):
        self.loop = None
        asyncio.set_event_loop_policy(None)

+    def test_async_gen_anext(self):
+        async def gen():
+            yield 1
+            yield 2
+        g = gen()
+        async def consume():
+            results = []
+            results.append(await anext(g))
+            results.append(await anext(g))
+            results.append(await anext(g, 'buckle my shoe'))
+            return results
+        res = self.loop.run_until_complete(consume())
+        self.assertEqual(res, [1, 2, 'buckle my shoe'])
+        with self.assertRaises(StopAsyncIteration):
+            self.loop.run_until_complete(consume())
+
+    def test_async_gen_aiter(self):
+        async def gen():
+            yield 1
+            yield 2
+        g = gen()
+        async def consume():
+            return [i async for i in aiter(g)]
+        res = self.loop.run_until_complete(consume())
+        self.assertEqual(res, [1, 2])
+
+    def test_async_gen_aiter_class(self):
+        results = []
+        class Gen:
+            async def __aiter__(self):
+                yield 1
+                yield 2
+        g = Gen()
+        async def consume():
+            ait = aiter(g)
+            while True:
+                try:
+                    results.append(await anext(ait))
+                except StopAsyncIteration:
+                    break
+        self.loop.run_until_complete(consume())
+        self.assertEqual(results, [1, 2])
+
+    def test_aiter_idempotent(self):
+        async def gen():
+            yield 1
+        applied_once = aiter(gen())
+        applied_twice = aiter(applied_once)
+        self.assertIs(applied_once, applied_twice)
+
+    def test_anext_bad_args(self):
+        async def gen():
+            yield 1
+        async def call_with_too_few_args():
+            await anext()
+        async def call_with_too_many_args():
+            await anext(gen(), 1, 3)
+        async def call_with_wrong_type_args():
+            await anext(1, gen())
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_too_few_args())
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_too_many_args())
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_wrong_type_args())
+
+    def test_aiter_bad_args(self):
+        async def gen():
+            yield 1
+        async def call_with_too_few_args():
+            await aiter()
+        async def call_with_too_many_args():
+            await aiter(gen(), 1)
+        async def call_with_wrong_type_arg():
+            await aiter(1)
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_too_few_args())
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_too_many_args())
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_wrong_type_arg())
+
    async def to_list(self, gen):
        res = []
        async for i in gen:
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@ -3860,6 +3860,9 @@ class TestSignatureDefinitions(unittest.TestCase):
        needs_groups = {"range", "slice", "dir", "getattr",
                        "next", "iter", "vars"}
        no_signature |= needs_groups
+        # These have unrepresentable parameter default values of NULL
+        needs_null = {"anext"}
+        no_signature |= needs_null
        # These need PEP 457 groups or a signature change to accept None
        needs_semantic_update = {"round"}
        no_signature |= needs_semantic_update
--- a/Misc/NEWS.d/next/Library/2018-08-24-01-08-09.bpo-31861.-q9RKJ.rst
+++ b/Misc/NEWS.d/next/Library/2018-08-24-01-08-09.bpo-31861.-q9RKJ.rst
@ -0,0 +1,2 @@
+Add builtins.aiter and builtins.anext.
+Patch by Joshua Bronson (@jab), Daniel Pope (@lordmauve), and Justin Wang (@justin39).
--- 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
--- a/Objects/iterobject.c
+++ b/Objects/iterobject.c
@ -157,7 +157,7 @@ PyTypeObject PySeqIter_Type = {
    PyObject_GenericGetAttr,                    /* tp_getattro */
    0,                                          /* tp_setattro */
    0,                                          /* tp_as_buffer */
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
    0,                                          /* tp_doc */
    (traverseproc)iter_traverse,                /* tp_traverse */
    0,                                          /* tp_clear */
@ -276,7 +276,7 @@ PyTypeObject PyCallIter_Type = {
    PyObject_GenericGetAttr,                    /* tp_getattro */
    0,                                          /* tp_setattro */
    0,                                          /* tp_as_buffer */
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
    0,                                          /* tp_doc */
    (traverseproc)calliter_traverse,            /* tp_traverse */
    0,                                          /* tp_clear */
@ -288,3 +288,91 @@ PyTypeObject PyCallIter_Type = {
 };


+/* -------------------------------------- */
+
+typedef struct {
+    PyObject_HEAD
+    PyObject *wrapped;
+    PyObject *default_value;
+} anextawaitableobject;
+
+static void
+anextawaitable_dealloc(anextawaitableobject *obj)
+{
+    _PyObject_GC_UNTRACK(obj);
+    Py_XDECREF(obj->wrapped);
+    Py_XDECREF(obj->default_value);
+    PyObject_GC_Del(obj);
+}
+
+static int
+anextawaitable_traverse(anextawaitableobject *obj, visitproc visit, void *arg)
+{
+    Py_VISIT(obj->wrapped);
+    Py_VISIT(obj->default_value);
+    return 0;
+}
+
+static PyObject *
+anextawaitable_iternext(anextawaitableobject *obj)
+{
+    PyObject *result = PyIter_Next(obj->wrapped);
+    if (result != NULL) {
+        return result;
+    }
+    if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration)) {
+        _PyGen_SetStopIterationValue(obj->default_value);
+    }
+    return NULL;
+}
+
+static PyAsyncMethods anextawaitable_as_async = {
+    PyObject_SelfIter,                          /* am_await */
+    0,                                          /* am_aiter */
+    0,                                          /* am_anext */
+    0,                                          /* am_send  */
+};
+
+PyTypeObject PyAnextAwaitable_Type = {
+    PyVarObject_HEAD_INIT(&PyType_Type, 0)
+    "anext_awaitable",                          /* tp_name */
+    sizeof(anextawaitableobject),               /* tp_basicsize */
+    0,                                          /* tp_itemsize */
+    /* methods */
+    (destructor)anextawaitable_dealloc,         /* tp_dealloc */
+    0,                                          /* tp_vectorcall_offset */
+    0,                                          /* tp_getattr */
+    0,                                          /* tp_setattr */
+    &anextawaitable_as_async,                   /* tp_as_async */
+    0,                                          /* tp_repr */
+    0,                                          /* tp_as_number */
+    0,                                          /* tp_as_sequence */
+    0,                                          /* tp_as_mapping */
+    0,                                          /* tp_hash */
+    0,                                          /* tp_call */
+    0,                                          /* tp_str */
+    PyObject_GenericGetAttr,                    /* tp_getattro */
+    0,                                          /* tp_setattro */
+    0,                                          /* tp_as_buffer */
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,    /* tp_flags */
+    0,                                          /* tp_doc */
+    (traverseproc)anextawaitable_traverse,      /* tp_traverse */
+    0,                                          /* tp_clear */
+    0,                                          /* tp_richcompare */
+    0,                                          /* tp_weaklistoffset */
+    PyObject_SelfIter,                          /* tp_iter */
+    (unaryfunc)anextawaitable_iternext,         /* tp_iternext */
+    0,                                          /* tp_methods */
+};
+
+PyObject *
+PyAnextAwaitable_New(PyObject *awaitable, PyObject *default_value)
+{
+    anextawaitableobject *anext = PyObject_GC_New(anextawaitableobject, &PyAnextAwaitable_Type);
+    Py_INCREF(awaitable);
+    anext->wrapped = awaitable;
+    Py_INCREF(default_value);
+    anext->default_value = default_value;
+    _PyObject_GC_TRACK(anext);
+    return (PyObject *)anext;
+}
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@ -1610,6 +1610,59 @@ supply its own iterator, or be a sequence.\n\
 In the second form, the callable is called until it returns the sentinel.");


+/*[clinic input]
+aiter as builtin_aiter
+
+    async_iterable: object
+    /
+
+Return an AsyncIterator for an AsyncIterable object.
+[clinic start generated code]*/
+
+static PyObject *
+builtin_aiter(PyObject *module, PyObject *async_iterable)
+/*[clinic end generated code: output=1bae108d86f7960e input=473993d0cacc7d23]*/
+{
+    return PyObject_GetAiter(async_iterable);
+}
+
+PyObject *PyAnextAwaitable_New(PyObject *, PyObject *);
+
+/*[clinic input]
+anext as builtin_anext
+
+    aiterator: object
+    default: object = NULL
+    /
+
+Return the next item from the async iterator.
+[clinic start generated code]*/
+
+static PyObject *
+builtin_anext_impl(PyObject *module, PyObject *aiterator,
+                   PyObject *default_value)
+/*[clinic end generated code: output=f02c060c163a81fa input=699d11f4e38eca24]*/
+{
+    PyTypeObject *t;
+    PyObject *awaitable;
+
+    t = Py_TYPE(aiterator);
+    if (t->tp_as_async == NULL || t->tp_as_async->am_anext == NULL) {
+        PyErr_Format(PyExc_TypeError,
+            "'%.200s' object is not an async iterator",
+            t->tp_name);
+        return NULL;
+    }
+
+    awaitable = (*t->tp_as_async->am_anext)(aiterator);
+    if (default_value == NULL) {
+        return awaitable;
+    }
+
+    return PyAnextAwaitable_New(awaitable, default_value);
+}
+
+
 /*[clinic input]
 len as builtin_len

@ -2890,11 +2943,13 @@ static PyMethodDef builtin_methods[] = {
    BUILTIN_ISINSTANCE_METHODDEF
    BUILTIN_ISSUBCLASS_METHODDEF
    {"iter",            (PyCFunction)(void(*)(void))builtin_iter,       METH_FASTCALL, iter_doc},
+    BUILTIN_AITER_METHODDEF
    BUILTIN_LEN_METHODDEF
    BUILTIN_LOCALS_METHODDEF
    {"max",             (PyCFunction)(void(*)(void))builtin_max,        METH_VARARGS | METH_KEYWORDS, max_doc},
    {"min",             (PyCFunction)(void(*)(void))builtin_min,        METH_VARARGS | METH_KEYWORDS, min_doc},
    {"next",            (PyCFunction)(void(*)(void))builtin_next,       METH_FASTCALL, next_doc},
+    BUILTIN_ANEXT_METHODDEF
    BUILTIN_OCT_METHODDEF
    BUILTIN_ORD_METHODDEF
    BUILTIN_POW_METHODDEF
--- a/Python/clinic/bltinmodule.c.h
+++ b/Python/clinic/bltinmodule.c.h
@ -530,6 +530,50 @@ PyDoc_STRVAR(builtin_hex__doc__,
 #define BUILTIN_HEX_METHODDEF    \
    {"hex", (PyCFunction)builtin_hex, METH_O, builtin_hex__doc__},

+PyDoc_STRVAR(builtin_aiter__doc__,
+"aiter($module, async_iterable, /)\n"
+"--\n"
+"\n"
+"Return an AsyncIterator for an AsyncIterable object.");
+
+#define BUILTIN_AITER_METHODDEF    \
+    {"aiter", (PyCFunction)builtin_aiter, METH_O, builtin_aiter__doc__},
+
+PyDoc_STRVAR(builtin_anext__doc__,
+"anext($module, aiterator, default=<unrepresentable>, /)\n"
+"--\n"
+"\n"
+"Return the next item from the async iterator.");
+
+#define BUILTIN_ANEXT_METHODDEF    \
+    {"anext", (PyCFunction)(void(*)(void))builtin_anext, METH_FASTCALL, builtin_anext__doc__},
+
+static PyObject *
+builtin_anext_impl(PyObject *module, PyObject *aiterator,
+                   PyObject *default_value);
+
+static PyObject *
+builtin_anext(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
+{
+    PyObject *return_value = NULL;
+    PyObject *aiterator;
+    PyObject *default_value = NULL;
+
+    if (!_PyArg_CheckPositional("anext", nargs, 1, 2)) {
+        goto exit;
+    }
+    aiterator = args[0];
+    if (nargs < 2) {
+        goto skip_optional;
+    }
+    default_value = args[1];
+skip_optional:
+    return_value = builtin_anext_impl(module, aiterator, default_value);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(builtin_len__doc__,
 "len($module, obj, /)\n"
 "--\n"
@ -830,4 +874,4 @@ builtin_issubclass(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
 exit:
    return return_value;
 }
-/*[clinic end generated code: output=e2fcf0201790367c input=a9049054013a1b77]*/
+/*[clinic end generated code: output=da9ae459e9233259 input=a9049054013a1b77]*/