gh-157044: Preserve asyncio call stacks through callable aiter()

Author: goutamadwant

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.

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
 

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])
 

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<worker>',
+                        ['a wait', 'a deep', 'a produce', 'a worker'],
+                    ])
+                    for name in ('deep', 'produce', 'worker'):
+                        self.assertIn(f'.<locals>.{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

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`.

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,
 };