From ba728ed30fcd65fc70ffabaa797a34a3dbdd535a Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Mon, 7 Sep 2026 14:42:18 -0700 Subject: [PATCH] gh-157044: Preserve asyncio call stacks through callable aiter() --- Doc/library/functions.rst | 7 ++ Lib/asyncio/graph.py | 3 + Lib/test/test_asyncgen.py | 81 +++++++++++++++++++ Lib/test/test_asyncio/test_graph.py | 45 +++++++++++ ...-09-07-12-00-00.gh-issue-157044.f8Ks2P.rst | 2 + Objects/iterobject.c | 21 ++++- 6 files changed, 158 insertions(+), 1 deletion(-) create mode 100644 Misc/NEWS.d/next/Library/2026-09-07-12-00-00.gh-issue-157044.f8Ks2P.rst diff --git a/Doc/library/functions.rst b/Doc/library/functions.rst index 013150535cb089c..c9952cf121312db 100644 --- a/Doc/library/functions.rst +++ b/Doc/library/functions.rst @@ -87,6 +87,13 @@ are always available. They are listed here in alphabetical order. The callable is only called when the result of :meth:`~object.__anext__` is awaited. + .. impl-detail:: + + The awaitable returned by :meth:`~object.__anext__` exposes the object + returned by *callable* through its read-only ``aw_wrapped`` attribute. + This attribute is ``None`` until the callable is invoked, and retains + the returned object after the awaitable completes or is closed. + *stop_exception* is an exception class or a tuple of exception classes. If *stop_value* is not specified, the iteration stops only when the callable raises an exception. diff --git a/Lib/asyncio/graph.py b/Lib/asyncio/graph.py index d5db59f5d6f5f36..529f9bf709bbfcf 100644 --- a/Lib/asyncio/graph.py +++ b/Lib/asyncio/graph.py @@ -65,6 +65,9 @@ def _build_graph_for_future( # A native async generator or duck-type compatible iterator st.append(FrameCallGraphEntry(coro.ag_frame)) coro = coro.ag_await + elif hasattr(coro, 'aw_wrapped'): + # An asynchronous callable iterator's frameless awaitable. + coro = coro.aw_wrapped else: break diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py index cdae58b3e89ae36..dc58dbd7542142c 100644 --- a/Lib/test/test_asyncgen.py +++ b/Lib/test/test_asyncgen.py @@ -909,12 +909,87 @@ async def spam(): def test_aiter_callable_awaitable(self): it = aiter(self.make_counter(), 10) awaitable = it.__anext__() + self.assertIsNone(awaitable.aw_wrapped) self.assertIsNone(awaitable.close()) + self.assertIsNone(awaitable.aw_wrapped) with self.assertRaises(RuntimeError): self.loop.run_until_complete(awaitable) awaitable = it.__anext__() with self.assertRaises(KeyError): awaitable.throw(KeyError('injected')) + self.assertIsNone(awaitable.aw_wrapped) + + def test_aiter_callable_wrapped(self): + async def produce(): + await awaitable() + await awaitable() + return 1 + + class CustomAwaitable: + def __init__(self): + self.iterator = self.iterate() + + def iterate(self): + yield ('result',) + + def __await__(self): + return self.iterator + + for factory in (produce, awaitable, CustomAwaitable): + for action in ('complete', 'close', 'throw'): + with self.subTest(factory=factory, action=action): + calls = [] + + def get_awaitable(): + wrapped = factory() + calls.append(wrapped) + return wrapped + + wrapper = anext(aiter(get_awaitable, object())) + try: + self.assertIsNone(wrapper.aw_wrapped) + self.assertEqual(calls, []) + with self.assertRaises(AttributeError): + wrapper.aw_wrapped = None + with self.assertRaises(AttributeError): + del wrapper.aw_wrapped + self.assertEqual(next(wrapper), ('result',)) + self.assertIs(wrapper.aw_wrapped, calls[0]) + if factory is produce: + delegate = calls[0].cr_await + self.assertEqual(wrapper.send(None), ('result',)) + self.assertIs(wrapper.aw_wrapped, calls[0]) + self.assertIsNot(calls[0].cr_await, delegate) + if action == 'complete': + with self.assertRaises(StopIteration): + wrapper.send(None) + elif action == 'throw': + with self.assertRaises(AwaitException): + wrapper.throw(AwaitException) + else: + wrapper.close() + self.assertIs(wrapper.aw_wrapped, calls[0]) + finally: + wrapper.close() + + def test_aiter_callable_wrapped_sentinel(self): + async def produce(): + return 1 + + iterator = aiter(produce, 1) + wrapper = anext(iterator) + with self.assertRaises(StopAsyncIteration): + wrapper.send(None) + wrapped = wrapper.aw_wrapped + self.assertIsNotNone(wrapped) + self.assertEqual(inspect.getcoroutinestate(wrapped), + inspect.CORO_CLOSED) + wrapper.close() + self.assertIs(wrapper.aw_wrapped, wrapped) + exhausted = anext(iterator) + with self.assertRaises(StopAsyncIteration): + exhausted.send(None) + self.assertIsNone(exhausted.aw_wrapped) def test_aiter_callable_cancel(self): # Cancellation is delivered to the awaited callable result @@ -931,9 +1006,15 @@ async def consume(): async def main(): task = asyncio.ensure_future(consume()) await asyncio.sleep(0) + wrapper = task.get_coro().cr_await + wrapped = wrapper.aw_wrapped + self.assertIsNotNone(wrapped) task.cancel() with self.assertRaises(asyncio.CancelledError): await task + self.assertIs(wrapper.aw_wrapped, wrapped) + self.assertEqual(inspect.getcoroutinestate(wrapped), + inspect.CORO_CLOSED) self.loop.run_until_complete(main()) self.assertEqual(cancelled, [1]) diff --git a/Lib/test/test_asyncio/test_graph.py b/Lib/test/test_asyncio/test_graph.py index 36841672e1f0f65..604a282ae3468fb 100644 --- a/Lib/test/test_asyncio/test_graph.py +++ b/Lib/test/test_asyncio/test_graph.py @@ -173,6 +173,51 @@ class FakeCoro: self.assertEqual(len(result.call_stack), 2) + async def test_stack_aiter_callable(self): + for nested in (False, True): + with self.subTest(nested=nested): + entered = asyncio.Event() + blocked = asyncio.Event() + + async def deep(): + entered.set() + await blocked.wait() + + async def produce(): + await deep() + return None + + async def worker(): + iterator = aiter(produce, None) + if nested: + iterator = aiter(iterator.__anext__, None) + async for _ in iterator: + pass + + task = asyncio.create_task(worker(), name='worker') + try: + await entered.wait() + stack, printed = capture_test_stack(fut=task) + self.assertEqual(stack[:2], [ + 'T', + ['a wait', 'a deep', 'a produce', 'a worker'], + ]) + for name in ('deep', 'produce', 'worker'): + self.assertIn(f'..{name}()', printed) + for limit, names in ( + (0, []), + (2, ['produce', 'worker']), + (-2, ['wait', 'deep']), + ): + graph = asyncio.capture_call_graph(task, limit=limit) + self.assertEqual( + [entry.frame.f_code.co_name + for entry in graph.call_stack], names) + finally: + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + async def test_stack_gather(self): stack_for_deep = None diff --git a/Misc/NEWS.d/next/Library/2026-09-07-12-00-00.gh-issue-157044.f8Ks2P.rst b/Misc/NEWS.d/next/Library/2026-09-07-12-00-00.gh-issue-157044.f8Ks2P.rst new file mode 100644 index 000000000000000..98a3170f041243b --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-09-07-12-00-00.gh-issue-157044.f8Ks2P.rst @@ -0,0 +1,2 @@ +Fix :func:`asyncio.print_call_graph` truncating the call stack at the awaitable +returned by the callable form of :func:`aiter`. diff --git a/Objects/iterobject.c b/Objects/iterobject.c index b5783c92c8eb689..add7fae9eeadc22 100644 --- a/Objects/iterobject.c +++ b/Objects/iterobject.c @@ -805,7 +805,7 @@ acallawaitable_start(acallawaitableobject *aw) } return -1; } - aw->aw_wrapped = awaitable; + FT_ATOMIC_STORE_PTR_RELEASE(aw->aw_wrapped, awaitable); return 0; } @@ -932,6 +932,24 @@ acallawaitable_close(PyObject *op, PyObject *Py_UNUSED(dummy)) return result; } +static PyObject * +acallawaitable_get_wrapped(PyObject *op, void *Py_UNUSED(closure)) +{ + acallawaitableobject *aw = acallawaitableobject_CAST(op); + PyObject *wrapped = FT_ATOMIC_LOAD_PTR_ACQUIRE(aw->aw_wrapped); + if (wrapped == NULL) { + Py_RETURN_NONE; + } + return Py_NewRef(wrapped); +} + +static PyGetSetDef acallawaitable_getset[] = { + {"aw_wrapped", acallawaitable_get_wrapped, NULL, + PyDoc_STR("Awaitable returned by the callable, or None before it is called."), + NULL}, + {NULL} +}; + static PyMethodDef acallawaitable_methods[] = { {"send", acallawaitable_send, METH_O, send_doc}, {"throw", acallawaitable_throw, METH_VARARGS, throw_doc}, @@ -958,4 +976,5 @@ PyTypeObject _PyACallIterAwaitable_Type = { .tp_iter = PyObject_SelfIter, .tp_iternext = acallawaitable_iternext, .tp_methods = acallawaitable_methods, + .tp_getset = acallawaitable_getset, };