bpo-31861: Add operator.aiter and operator.anext

python · jab · Aug 24, 2018 · Sep 7, 2018 · Nov 30, 2020 · Dec 4, 2020
commit 44b2a6032d6edb82a79daeccec6994daa8a601e8
diff --git a/Lib/operator.py b/Lib/operator.py
@@ -10,16 +10,18 @@
 This is the pure Python implementation of the module.
 """
 
-__all__ = ['abs', 'add', 'and_', 'attrgetter', 'concat', 'contains', 'countOf',
-           'delitem', 'eq', 'floordiv', 'ge', 'getitem', 'gt', 'iadd', 'iand',
-           'iconcat', 'ifloordiv', 'ilshift', 'imatmul', 'imod', 'imul',
-           'index', 'indexOf', 'inv', 'invert', 'ior', 'ipow', 'irshift',
-           'is_', 'is_not', 'isub', 'itemgetter', 'itruediv', 'ixor', 'le',
-           'length_hint', 'lshift', 'lt', 'matmul', 'methodcaller', 'mod',
-           'mul', 'ne', 'neg', 'not_', 'or_', 'pos', 'pow', 'rshift',
-           'setitem', 'sub', 'truediv', 'truth', 'xor']
+__all__ = [
+    'abs', 'add', 'aiter', 'anext', 'and_', 'attrgetter', 'concat', 'contains',
+    'countOf', 'delitem', 'eq', 'floordiv', 'ge', 'getitem', 'gt', 'iadd',
+    'iand', 'iconcat', 'ifloordiv', 'ilshift', 'imatmul', 'imod', 'imul',
+    'index', 'indexOf', 'inv', 'invert', 'ior', 'ipow', 'irshift', 'is_',
+    'is_not', 'isub', 'itemgetter', 'itruediv', 'ixor', 'le', 'length_hint',
+    'lshift', 'lt', 'matmul', 'methodcaller', 'mod', 'mul', 'ne', 'neg', 'not_',
+    'or_', 'pos', 'pow', 'rshift', 'setitem', 'sub', 'truediv', 'truth', 'xor',
+]
 
 from builtins import abs as _abs
+from collections.abc import AsyncIterable, AsyncIterator
 
 
 # Comparison Operations *******************************************************#
@@ -404,6 +406,55 @@ def ixor(a, b):
     return a
 
 
+# Asynchronous Iterator Operations ********************************************#
+
+async def aiter(*args):
+    """aiter(async_iterable) -> async_iterator
+    aiter(async_callable, sentinel) -> async_iterator
+
+    An async version of the iter() builtin.
+    """
+    lenargs = len(args)
+    if lenargs != 1 and lenargs != 2:
+        raise TypeError(f'aiter expected 1 or 2 arguments, got {lenargs}')
+    if lenargs == 1:
+        obj, = args
+        if not isinstance(obj, AsyncIterable):
+            raise TypeError(f'aiter expected an AsyncIterable, got {type(obj)}')
+        async for i in obj.__aiter__():
+            yield i
+        return
+    # lenargs == 2
+    async_callable, sentinel = args
+    while True:
+        value = await async_callable()
+        if value == sentinel:
+            break
+        yield value
+
+
+async def anext(*args):
+    """anext(async_iterator[, default])
+
+    Return the next item from the async iterator.
+    If default is given and the iterator is exhausted,
+    it is returned instead of raising StopAsyncIteration.
+    """
+    lenargs = len(args)
+    if lenargs != 1 and lenargs != 2:
+        raise TypeError(f'anext expected 1 or 2 arguments, got {lenargs}')
+    ait = args[0]
+    if not isinstance(ait, AsyncIterator):
+        raise TypeError(f'anext expected an AsyncIterable, got {type(ait)}')
+    anxt = ait.__anext__
+    try:
+        return await anxt()
+    except StopAsyncIteration:
+        if lenargs == 1:
+            raise
+        return args[1]  # default
+
+
 try:
     from _operator import *
 except ImportError:

diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py
@@ -1,4 +1,5 @@
 import inspect
+import operator
 import types
 import unittest
 
@@ -372,6 +373,80 @@ def tearDown(self):
         self.loop = None
         asyncio.set_event_loop_policy(None)
 
+    def test_async_gen_operator_anext(self):
+        async def gen():
+            yield 1
+            yield 2
+        g = gen()
+        async def consume():
+            results = []
+            results.append(await operator.anext(g))
+            results.append(await operator.anext(g))
+            results.append(await operator.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_operator_aiter(self):
+        async def gen():
+            yield 1
+            yield 2
+        g = gen()
+        async def consume():
+            return [i async for i in operator.aiter(g)]
+        res = self.loop.run_until_complete(consume())
+        self.assertEqual(res, [1, 2])
+
+    def test_async_gen_operator_aiter_class(self):
+        loop = self.loop
+        class Gen:
+            async def __aiter__(self):
+                yield 1
+                await asyncio.sleep(0.01, loop=loop)
+                yield 2
+        g = Gen()
+        async def consume():
+            return [i async for i in operator.aiter(g)]
+        res = self.loop.run_until_complete(consume())
+        self.assertEqual(res, [1, 2])
+
+    def test_async_gen_operator_aiter_2_arg(self):
+        async def gen():
+            yield 1
+            yield 2
+            yield None
+        g = gen()
+        async def foo():
+            return await operator.anext(g)
+        async def consume():
+            return [i async for i in operator.aiter(foo, None)]
+        res = self.loop.run_until_complete(consume())
+        self.assertEqual(res, [1, 2])
+
+    def test_operator_anext_bad_args(self):
+        self._test_bad_args(operator.anext)
+
+    def test_operator_aiter_bad_args(self):
+        self._test_bad_args(operator.aiter)
+
+    def _test_bad_args(self, afn):
+        async def gen():
+            yield 1
+        async def call_with_no_args():
+            await afn()
+        async def call_with_3_args():
+            await afn(gen(), 1, 2)
+        async def call_with_bad_args():
+            await afn(1, gen())
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_no_args())
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_3_args())
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(call_with_bad_args())
+
     async def to_list(self, gen):
         res = []
         async for i in gen:

diff --git 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 @@
+Add the operator.aiter and operator.anext functions. Patch by Josh Bronson.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Add the operator.aiter and operator.anext functions. Patch by Josh Bronson.