Implement, test, and document asyncio.print_call_stack()

1st1 · 1st1 · commit 90a54d95dc85 · 2024-09-28T15:04:11.000-07:00
diff --git a/Doc/library/asyncio-stack.rst b/Doc/library/asyncio-stack.rst
@@ -23,9 +23,9 @@ a suspended *future*.
    Capture the async call stack for the current task or the provided
    :class:`Task` or :class:`Future`.
 
-   The function receives an optional keyword-only *future* argument.
-   If not passed, the current task will be used. If there's no current task,
-   the function returns ``None``.
+   The function recieves an optional keyword-only *future* argument.
+   If not passed, the current running task will be used. If there's no
+   current task, the function returns ``None``.
 
    Returns a ``FutureCallStack`` named tuple:
 
@@ -49,6 +49,17 @@ a suspended *future*.
       Where ``coroutine`` is a coroutine object of an awaiting coroutine
       or asyncronous generator.
 
+.. function:: print_call_stack(*, future=None, file=None)
+
+   Print the async call stack for the current task or the provided
+   :class:`Task` or :class:`Future`.
+
+   The function recieves an optional keyword-only *future* argument.
+   If not passed, the current running task will be used. If there's no
+   current task, the function returns ``None``.
+
+   If *file* is not specified the function will print to :data:`sys.stdout`.
+
 
 Low level utility functions
 ===========================
@@ -70,6 +81,10 @@ the tasks they wrap or control.
    :class:`asyncio.Future <Future>` or :class:`asyncio.Task <Task>` or
    their subclasses, otherwise the call would have no effect.
 
+   A call to ``future_add_to_awaited_by()`` must be followed by an
+   eventual call to the ``future_discard_from_awaited_by()`` function
+   with the same arguments.
+
 
 .. function:: future_discard_from_awaited_by(future, waiter, /)
 
diff --git a/Lib/asyncio/stack.py b/Lib/asyncio/stack.py
@@ -10,6 +10,7 @@
 
 __all__ = (
     'capture_call_stack',
+    'print_call_stack',
     'FrameCallStackEntry',
     'CoroutineCallStackEntry',
     'FutureCallStack',
@@ -110,15 +111,15 @@ def capture_call_stack(*, future: any = None) -> FutureCallStack | None:
         # Check if we're in a context of a running event loop;
         # if yes - check if the passed future is the currently
         # running task or not.
-        if loop is None or future is not tasks.current_task():
+        if loop is None or future is not tasks.current_task(loop=loop):
             return _build_stack_for_future(future)
         # else: future is the current task, move on.
     else:
         if loop is None:
             raise RuntimeError(
                 'capture_call_stack() is called outside of a running '
                 'event loop and no *future* to introspect was provided')
-        future = tasks.current_task()
+        future = tasks.current_task(loop=loop)
 
     if future is None:
         # This isn't a generic call stack introspection utility. If we
@@ -158,3 +159,76 @@ def capture_call_stack(*, future: any = None) -> FutureCallStack | None:
             awaited_by.append(_build_stack_for_future(parent))
 
     return FutureCallStack(future, call_stack, awaited_by)
+
+
+def print_call_stack(*, future: any = None, file=None) -> None:
+    """Print async call stack for the current task or the provided Future."""
+
+    stack = capture_call_stack(future=future)
+    if stack is None:
+        return
+
+    buf = []
+
+    def render_level(st: FutureCallStack, level: int = 0):
+        def add_line(line: str):
+            buf.append(level * '    ' + line)
+
+        if isinstance(st.future, tasks.Task):
+            add_line(
+                f'* Task(name={st.future.get_name()!r}, id=0x{id(st.future):x})'
+            )
+        else:
+            add_line(
+                f'* Future(id=0x{id(st.future):x})'
+            )
+
+        if st.call_stack:
+            add_line(
+                f'  + Call stack:'
+            )
+            for ste in st.call_stack:
+                if isinstance(ste, FrameCallStackEntry):
+                    f = ste.frame
+                    add_line(
+                        f'  | * {f.f_code.co_qualname}()'
+                    )
+                    add_line(
+                        f'  |   {f.f_code.co_filename}:{f.f_lineno}'
+                    )
+                else:
+                    assert isinstance(ste, CoroutineCallStackEntry)
+                    c = ste.coroutine
+
+                    try:
+                        f = c.cr_frame
+                        code = c.cr_code
+                        tag = 'async'
+                    except AttributeError:
+                        try:
+                            f = c.ag_frame
+                            code = c.ag_code
+                            tag = 'async generator'
+                        except AttributeError:
+                            f = c.gi_frame
+                            code = c.gi_code
+                            tag = 'generator'
+
+                    add_line(
+                        f'  | * {tag} {code.co_qualname}()'
+                    )
+                    add_line(
+                        f'  |   {f.f_code.co_filename}:{f.f_lineno}'
+                    )
+
+        if st.awaited_by:
+            add_line(
+                f'  + Awaited by:'
+            )
+            for fut in st.awaited_by:
+                render_level(fut, level + 1)
+
+    render_level(stack)
+    rendered = '\n'.join(buf)
+
+    print(rendered, file=file)
diff --git a/Lib/test/test_asyncio/test_stack.py b/Lib/test/test_asyncio/test_stack.py
@@ -1,4 +1,5 @@
 import asyncio
+import io
 import unittest
 
 
@@ -37,8 +38,11 @@ def walk(s):
 
         return ret
 
+    buf = io.StringIO()
+    asyncio.print_call_stack(future=fut, file=buf)
+
     stack = asyncio.capture_call_stack(future=fut)
-    return walk(stack)
+    return walk(stack), buf.getvalue()
 
 
 class TestCallStack(unittest.IsolatedAsyncioTestCase):
@@ -73,7 +77,7 @@ async def main():
 
         await main()
 
-        self.assertEqual(stack_for_c5, [
+        self.assertEqual(stack_for_c5[0], [
             # task name
             'T<c2_root>',
             # call stack
@@ -102,6 +106,11 @@ async def main():
             ]
         ])
 
+        self.assertIn(
+            '* async TestCallStack.test_stack_tgroup()',
+            stack_for_c5[1])
+
+
     async def test_stack_async_gen(self):
 
         stack_for_gen_nested_call = None
@@ -122,7 +131,7 @@ async def main():
 
         await main()
 
-        self.assertEqual(stack_for_gen_nested_call, [
+        self.assertEqual(stack_for_gen_nested_call[0], [
             'T<anon>',
             [
                 's capture_test_stack',
@@ -134,6 +143,10 @@ async def main():
             []
         ])
 
+        self.assertIn(
+            'async generator TestCallStack.test_stack_async_gen.<locals>.gen()',
+            stack_for_gen_nested_call[1])
+
     async def test_stack_gather(self):
 
         stack_for_deep = None
@@ -155,7 +168,7 @@ async def main():
 
         await main()
 
-        self.assertEqual(stack_for_deep, [
+        self.assertEqual(stack_for_deep[0], [
             'T<anon>',
             ['s capture_test_stack', 'a deep', 'a c1'],
             [
@@ -181,7 +194,7 @@ async def main():
 
         await main()
 
-        self.assertEqual(stack_for_shield, [
+        self.assertEqual(stack_for_shield[0], [
             'T<anon>',
             ['s capture_test_stack', 'a deep', 'a c1'],
             [
@@ -208,7 +221,7 @@ async def main():
 
         await main()
 
-        self.assertEqual(stack_for_inner, [
+        self.assertEqual(stack_for_inner[0], [
             'T<anon>',
             ['s capture_test_stack', 'a inner', 'a c1'],
             [
@@ -248,7 +261,7 @@ async def main(t1, t2):
             await t1
             await t2
 
-        self.assertEqual(stack_for_inner, [
+        self.assertEqual(stack_for_inner[0], [
             'T<anon>',
             ['s capture_test_stack', 'a inner', 'a c1'],
             [
@@ -279,7 +292,7 @@ async def main():
 
         await main()
 
-        self.assertEqual(stack_for_inner, [
+        self.assertEqual(stack_for_inner[0], [
             'T<there there>',
             ['s capture_test_stack', 'a inner', 'a c1'],
             [['T<anon>', ['a c2', 'a main', 'a test_stack_task'], []]]
@@ -316,7 +329,7 @@ async def main():
 
         await main()
 
-        self.assertEqual(stack_for_fut,
+        self.assertEqual(stack_for_fut[0],
             ['F',
             [],
             [
@@ -330,3 +343,5 @@ async def main():
                 ],
             ]]
         )
+
+        self.assertTrue(stack_for_fut[1].startswith('* Future(id='))
