Skip to content

Commit cb12c5e

Browse files
gh-157361: Implement aiter() in Python instead of C (#157687)
gh-157361: Implement aiter() in Python instead of C Move both forms of aiter() into the frozen _pybuiltins module and remove the C async_callable_iterator types. The callable form now returns an iterator whose __anext__() is a plain coroutine, so call-graph tools can see through it into the awaited result of the callable.
1 parent 4c73f50 commit cb12c5e

11 files changed

Lines changed: 211 additions & 511 deletions

File tree

‎Include/internal/pycore_interp_structs.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -538,7 +538,7 @@ struct _py_func_state {
538538
If you add a new static type to the standard library, you may have to
539539
update one of these numbers.
540540
*/
541-
#define _Py_NUM_MANAGED_PREINITIALIZED_TYPES 122
541+
#define _Py_NUM_MANAGED_PREINITIALIZED_TYPES 120
542542
#define _Py_MAX_MANAGED_STATIC_BUILTIN_TYPES \
543543
(_Py_NUM_MANAGED_PREINITIALIZED_TYPES + 83)
544544
#define _Py_MAX_MANAGED_STATIC_EXT_TYPES 10

‎Include/internal/pycore_iterobject.h‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,20 +8,12 @@ extern "C" {
88
# error "this header requires Py_BUILD_CORE define"
99
#endif
1010

11-
extern PyTypeObject _PyACallIter_Type;
12-
extern PyTypeObject _PyACallIterAwaitable_Type;
13-
1411
// Like PyCallIter_New(), but the iteration also stops when *callable* raises
1512
// an exception matching *stop_exc* (an exception class or a tuple of exception
1613
// classes). *sentinel* can be NULL; NULL *stop_exc* means StopIteration.
1714
extern PyObject *_PyCallIter_NewEx(PyObject *callable, PyObject *sentinel,
1815
PyObject *stop_exc);
1916

20-
// The asynchronous counterpart of _PyCallIter_NewEx(): the result of
21-
// *callable* is awaited, and NULL *stop_exc* means StopAsyncIteration.
22-
extern PyObject *_PyACallIter_New(PyObject *callable, PyObject *sentinel,
23-
PyObject *stop_exc);
24-
2517
#ifdef __cplusplus
2618
}
2719
#endif

‎Lib/_pybuiltins.py‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,96 @@
55
copied into the ``builtins`` module.
66
"""
77

8-
__all__ = ['anext']
8+
__all__ = ['aiter', 'anext']
99

1010
_NOT_GIVEN = sentinel("_NOT_GIVEN")
1111

1212

13+
def aiter(object, /, stop_value=_NOT_GIVEN, *, stop_exception=_NOT_GIVEN):
14+
"""Return an AsyncIterator for an AsyncIterable object.
15+
16+
In the second form, the callable is called and its result is awaited
17+
until it returns the stop value or raises the specified exception.
18+
"""
19+
if stop_value is _NOT_GIVEN and stop_exception is _NOT_GIVEN:
20+
cls = type(object)
21+
try:
22+
# Looked up on the type, like the C slot am_aiter.
23+
aiter_method = cls.__aiter__
24+
except AttributeError:
25+
raise TypeError(
26+
f"'{cls.__name__}' object is not an async iterable"
27+
) from None
28+
iterator = aiter_method(object)
29+
if not hasattr(type(iterator), '__anext__'):
30+
raise TypeError(
31+
f"{cls.__name__}.__aiter__() must return an async iterator, "
32+
f"not {type(iterator).__name__}"
33+
)
34+
return iterator
35+
36+
if not callable(object):
37+
raise TypeError("aiter(): the first argument must be callable")
38+
if stop_exception is _NOT_GIVEN:
39+
stop_exception = StopAsyncIteration
40+
else:
41+
# Like _PyEval_CheckExceptTypeValid(): an exception class or a tuple
42+
# of exception classes.
43+
if isinstance(stop_exception, tuple):
44+
items = stop_exception
45+
else:
46+
items = (stop_exception,)
47+
for item in items:
48+
if not (isinstance(item, type) and issubclass(item, BaseException)):
49+
raise TypeError(
50+
"catching classes that do not inherit from BaseException "
51+
"is not allowed"
52+
)
53+
return async_callable_iterator(object, stop_value, stop_exception)
54+
55+
56+
class async_callable_iterator:
57+
"""The asynchronous counterpart of callable_iterator.
58+
59+
The callable is called and its result is awaited for every
60+
__anext__().
61+
"""
62+
63+
__slots__ = ('_callable', '_stop_value', '_stop_exception')
64+
__module__ = 'builtins'
65+
66+
def __init__(self, callable, stop_value, stop_exception):
67+
# _callable is set to None when the iterator is exhausted.
68+
self._callable = callable
69+
self._stop_value = stop_value
70+
self._stop_exception = stop_exception
71+
72+
def __aiter__(self):
73+
return self
74+
75+
async def __anext__(self):
76+
callable = self._callable
77+
if callable is None:
78+
raise StopAsyncIteration
79+
try:
80+
value = await callable()
81+
except self._stop_exception:
82+
self._exhaust()
83+
raise StopAsyncIteration from None
84+
except (StopIteration, StopAsyncIteration) as exc:
85+
raise RuntimeError(
86+
f"callable raised {type(exc).__name__}") from exc
87+
stop_value = self._stop_value
88+
if stop_value is not _NOT_GIVEN and stop_value == value:
89+
self._exhaust()
90+
raise StopAsyncIteration
91+
return value
92+
93+
def _exhaust(self):
94+
self._callable = None
95+
self._stop_value = _NOT_GIVEN
96+
97+
1398
def anext(async_iterator, default=_NOT_GIVEN, /):
1499
"""Return the next item from the async iterator.
15100

‎Lib/test/test_asyncio/test_graph.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,45 @@ async def main():
148148
'async generator CallStackTestBase.test_stack_async_gen.<locals>.gen()',
149149
stack_for_gen_nested_call[1])
150150

151+
async def test_stack_anext(self):
152+
# anext() without a default returns the __anext__() coroutine
153+
# itself, so nothing is inserted between it and main().
154+
155+
loop = asyncio.get_running_loop()
156+
blocker = loop.create_future()
157+
158+
async def inner():
159+
await blocker
160+
161+
class AIter:
162+
def __aiter__(self):
163+
return self
164+
165+
async def __anext__(self):
166+
await inner()
167+
return 1
168+
169+
async def main():
170+
await anext(AIter())
171+
172+
task = asyncio.create_task(main(), name='anext task')
173+
await asyncio.sleep(0)
174+
try:
175+
stack = capture_test_stack(fut=task)
176+
finally:
177+
blocker.set_result(None)
178+
await task
179+
180+
self.assertEqual(stack[0], [
181+
'T<anext task>',
182+
[
183+
'a inner',
184+
'a __anext__',
185+
'a main',
186+
],
187+
[]
188+
])
189+
151190
async def test_stack_anext_default(self):
152191
# anext() with a default wraps the awaitable in a coroutine, so the
153192
# call graph of a suspended task sees through it into __anext__().
@@ -188,6 +227,84 @@ async def main():
188227
[]
189228
])
190229

230+
async def test_stack_aiter(self):
231+
# The one-argument form of aiter() returns the object's own
232+
# iterator, so it adds no frame of its own to the call graph.
233+
234+
loop = asyncio.get_running_loop()
235+
blocker = loop.create_future()
236+
237+
async def inner():
238+
await blocker
239+
240+
class AIter:
241+
def __aiter__(self):
242+
return self
243+
244+
async def __anext__(self):
245+
await inner()
246+
return 1
247+
248+
async def main():
249+
async for _ in aiter(AIter()):
250+
break
251+
252+
task = asyncio.create_task(main(), name='aiter task')
253+
await asyncio.sleep(0)
254+
try:
255+
stack = capture_test_stack(fut=task)
256+
finally:
257+
blocker.set_result(None)
258+
await task
259+
260+
self.assertEqual(stack[0], [
261+
'T<aiter task>',
262+
[
263+
'a inner',
264+
'a __anext__',
265+
'a main',
266+
],
267+
[]
268+
])
269+
270+
async def test_stack_aiter_callable(self):
271+
# aiter(callable, ...) returns an iterator whose __anext__() is a
272+
# coroutine, so the call graph of a suspended task sees through it
273+
# into the awaited result of the callable.
274+
275+
loop = asyncio.get_running_loop()
276+
blocker = loop.create_future()
277+
278+
async def inner():
279+
await blocker
280+
281+
async def spam():
282+
await inner()
283+
return 1
284+
285+
async def main():
286+
async for _ in aiter(spam, 1):
287+
pass
288+
289+
task = asyncio.create_task(main(), name='aiter task')
290+
await asyncio.sleep(0)
291+
try:
292+
stack = capture_test_stack(fut=task)
293+
finally:
294+
blocker.set_result(None)
295+
await task
296+
297+
self.assertEqual(stack[0], [
298+
'T<aiter task>',
299+
[
300+
'a inner',
301+
'a spam',
302+
'a __anext__',
303+
'a main',
304+
],
305+
[]
306+
])
307+
191308
def test_ag_frame_used_for_async_generator(self):
192309
# Regression test for gh-148736: the ag_await branch of
193310
# _build_graph_for_future must read ag_frame, not cr_frame.

‎Lib/test/test_inspect/test_inspect.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6229,7 +6229,7 @@ def test_builtins_have_signatures(self):
62296229
needs_groups = {"range", "slice", "getattr", "next"}
62306230
no_signature |= needs_groups
62316231
# These have unrepresentable parameter default values of NULL
6232-
unsupported_signature = {"aiter", "iter", "dir", "vars"}
6232+
unsupported_signature = {"iter", "dir", "vars"}
62336233
# These need *args support in Argument Clinic
62346234
needs_varargs = {"min", "max", "__build_class__"}
62356235
no_signature |= needs_varargs
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Implement :func:`aiter` in Python instead of C, in the frozen ``_pybuiltins``
2+
module. The async iterator returned by ``aiter(callable, ...)`` now has a
3+
plain coroutine ``__anext__()``, so introspection tools such as
4+
:func:`asyncio.print_call_graph` can see through it into the awaited
5+
callable result.

0 commit comments

Comments
 (0)