gh-160804: Support generator based awaitables in asyncio introspection (#160805)

Author: kumaraditya303

Lib/asyncio/graph.py

@@ -65,6 +65,11 @@ 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, 'gi_yieldfrom'):
+            # A generator, e.g. returned by __await__ or @types.coroutine
+            if coro.gi_frame is not None:
+                st.append(FrameCallGraphEntry(coro.gi_frame))
+            coro = coro.gi_yieldfrom
         else:
             break
 

Lib/test/test_asyncio/test_graph.py

@@ -1,6 +1,7 @@
 import asyncio
 import io
 import sys
+import types
 import unittest
 from unittest import mock
 
@@ -11,6 +12,7 @@ def tearDownModule():
 
 
 def capture_test_stack(*, fut=None, depth=1):
+    future_await_code = asyncio.futures._PyFuture.__await__.__code__
 
     def walk(s):
         ret = [
@@ -27,8 +29,12 @@ def walk(s):
                             f"a {entry.frame.f_generator.cr_code.co_name}"
                             if hasattr(entry.frame.f_generator, 'cr_code') else
                             f"ag {entry.frame.f_generator.ag_code.co_name}"
+                            if hasattr(entry.frame.f_generator, 'ag_code') else
+                            f"g {entry.frame.f_generator.gi_code.co_name}"
                         )
                 ) for entry in s.call_stack
+                # Skip the pure Python Future.__await__ generator frame
+                if entry.frame.f_code is not future_await_code
             ]
         )
 
@@ -305,6 +311,73 @@ async def main():
             []
         ])
 
+    async def test_stack_await_generator(self):
+        # The call graph of a suspended task sees through generators
+        # returned by __await__() and the generators they yield from.
+
+        loop = asyncio.get_running_loop()
+        blocker = loop.create_future()
+
+        def wait():
+            yield from blocker
+
+        class Awaitable:
+            def __await__(self):
+                return (yield from wait())
+
+        async def main():
+            await Awaitable()
+
+        task = asyncio.create_task(main(), name='await task')
+        await asyncio.sleep(0)
+        try:
+            stack = capture_test_stack(fut=task)
+        finally:
+            blocker.set_result(None)
+            await task
+
+        self.assertEqual(stack[0], [
+            'T<await task>',
+            [
+                'g wait',
+                'g __await__',
+                'a main',
+            ],
+            []
+        ])
+
+    async def test_stack_types_coroutine(self):
+        loop = asyncio.get_running_loop()
+        blocker = loop.create_future()
+
+        async def inner():
+            await blocker
+
+        @types.coroutine
+        def legacy():
+            yield from inner()
+
+        async def main():
+            await legacy()
+
+        task = asyncio.create_task(main(), name='legacy task')
+        await asyncio.sleep(0)
+        try:
+            stack = capture_test_stack(fut=task)
+        finally:
+            blocker.set_result(None)
+            await task
+
+        self.assertEqual(stack[0], [
+            'T<legacy task>',
+            [
+                'a inner',
+                'g legacy',
+                'a main',
+            ],
+            []
+        ])
+
     def test_ag_frame_used_for_async_generator(self):
         # Regression test for gh-148736: the ag_await branch of
         # _build_graph_for_future must read ag_frame, not cr_frame.

Lib/test/test_external_inspection.py

@@ -8,6 +8,7 @@
 import socket
 import threading
 import time
+import types
 from contextlib import contextmanager
 from asyncio import staggered, taskgroups, base_events, tasks
 from unittest.mock import ANY
@@ -586,6 +587,54 @@ async def main():
         stacks = self._extract_coroutine_funcnames(asyncio.run(main()))
         self.assertEqual(stacks["anext"], ["agen"])
 
+    @skip_if_not_supported
+    def test_generator_awaitable_stack_is_not_truncated(self):
+        def wait():
+            yield from asyncio.get_running_loop().create_future()
+
+        class Awaitable:
+            def __await__(self):
+                return (yield from wait())
+
+        async def inner():
+            await asyncio.sleep(10_000)
+
+        @types.coroutine
+        def legacy():
+            yield from inner()
+
+        async def consumer(aw):
+            await aw
+
+        async def main():
+            await_task = asyncio.create_task(
+                consumer(Awaitable()), name="__await__")
+            legacy_task = asyncio.create_task(
+                consumer(legacy()), name="types.coroutine")
+            await asyncio.sleep(0)
+            return RemoteUnwinder(os.getpid()).get_all_awaited_by()
+
+        stacks = self._extract_coroutine_funcnames(asyncio.run(main()))
+        self.assertEqual(stacks["__await__"], ["wait", "__await__", "consumer"])
+        self.assertEqual(stacks["types.coroutine"],
+                         ["sleep", "inner", "legacy", "consumer"])
+
+    @skip_if_not_supported
+    def test_generator_awaitable_running_task_stack(self):
+        class Awaitable:
+            def __await__(self):
+                return RemoteUnwinder(os.getpid()).get_async_stack_trace()
+                yield
+
+        async def consumer():
+            return await Awaitable()
+
+        async def main():
+            return await asyncio.create_task(consumer(), name="consumer")
+
+        stacks = self._extract_coroutine_funcnames(asyncio.run(main()))
+        self.assertEqual(stacks["consumer"], ["__await__", "consumer"])
+
     @skip_if_not_supported
     def test_class_named_like_async_generator_wrapper_is_not_walked(self):
         # gh-159620

Misc/NEWS.d/next/Library/2026-10-28-07-26-47.gh-issue-160804.R5HeTQ.rst

@@ -0,0 +1,3 @@
+:mod:`asyncio` call graph introspection now includes frames of generator
+based awaitables, such as generators returned by ``__await__`` and
+:func:`types.coroutine` generators.

Modules/_remote_debugging/asyncio.c

@@ -241,8 +241,8 @@ parse_task_name(
  * ============================================================================ */
 
 // gh-159620
-/* Return in *coro_addr the coroutine or async generator to walk for the
-   awaitable at addr, or 0 if it is neither (for example, a future). */
+/* Return in *coro_addr the coroutine, generator or async generator to walk
+   for the awaitable at addr, or 0 if it is none (for example, a future). */
 static int
 get_coro_like_address(
     RemoteUnwinderObject *unwinder,
@@ -279,7 +279,7 @@ get_coro_like_address(
         return -1;
     }
 
-    if (strcmp(name, "coroutine") == 0) {
+    if (strcmp(name, "coroutine") == 0 || strcmp(name, "generator") == 0) {
         *coro_addr = addr;
     }
     else if (strcmp(name, "async_generator_asend") == 0