GH-91048: Add utils for printing the call stack for asyncio tasks

python · pablogsal · May 4, 2025 · May 1, 2025 · May 2, 2025 · May 2, 2025
commit b11469a4d20877e07f7e11b2ee96d596692eb5d4
@@ -0,0 +1,145 @@
+import argparse
+from collections import defaultdict
+from itertools import count
+from enum import Enum
+import sys
+from _remotedebugging import get_all_awaited_by
+
+
+class NodeType(Enum):
+    COROUTINE = 1
+    TASK = 2
+
+
+# ─── indexing helpers ───────────────────────────────────────────
+def _index(result):
+    id2name, awaits = {}, []
+    for _thr_id, tasks in result:
+        for tid, tname, awaited in tasks:
+            id2name[tid] = tname
+            for stack, parent_id in awaited:
+                awaits.append((parent_id, stack, tid))
+    return id2name, awaits
+
+
+def _build_tree(id2name, awaits):
+    id2label = {(NodeType.TASK, tid): name for tid, name in id2name.items()}
+    children = defaultdict(list)
+    cor_names = defaultdict(dict)  # (parent) -> {frame: node}
+    cor_id_seq = count(1)
+
+    def _cor_node(parent_key, frame_name):
+        """Return an existing or new (NodeType.COROUTINE, …) node under *parent_key*."""
+        bucket = cor_names[parent_key]
+        if frame_name in bucket:
+            return bucket[frame_name]
+        node_key = (NodeType.COROUTINE, f"c{next(cor_id_seq)}")
+        id2label[node_key] = frame_name
+        children[parent_key].append(node_key)
+        bucket[frame_name] = node_key
+        return node_key
+
+    # touch every task so it’s present even if it awaits nobody
+    for tid in id2name:
+        children[(NodeType.TASK, tid)]
+
+    # lay down parent ➜ …frames… ➜ child paths
+    for parent_id, stack, child_id in awaits:
+        cur = (NodeType.TASK, parent_id)
+        for frame in reversed(stack):  # outer-most → inner-most
+            cur = _cor_node(cur, frame)
+        child_key = (NodeType.TASK, child_id)
+        if child_key not in children[cur]:
+            children[cur].append(child_key)
+
+    return id2label, children
+
+
+def _roots(id2label, children):
+    roots = [n for n, lbl in id2label.items() if lbl == "Task-1"]
+    if roots:
+        return roots
+    all_children = {c for kids in children.values() for c in kids}
+    return [n for n in id2label if n not in all_children]
+
+
+# ─── PRINT TREE FUNCTION ───────────────────────────────────────
+def print_async_tree(result, task_emoji="(T)", cor_emoji="", printer=print):
+    """
+    Pretty-print the async call tree produced by `get_all_async_stacks()`,
+    prefixing tasks with *task_emoji* and coroutine frames with *cor_emoji*.
+    """
+    id2name, awaits = _index(result)
+    labels, children = _build_tree(id2name, awaits)
+
+    def pretty(node):
+        flag = task_emoji if node[0] == NodeType.TASK else cor_emoji
+        return f"{flag} {labels[node]}"
+
+    def render(node, prefix="", last=True, buf=None):
+        if buf is None:
+            buf = []
+        buf.append(f"{prefix}{'└── ' if last else '├── '}{pretty(node)}")
+        new_pref = prefix + ("    " if last else "│   ")
+        kids = children.get(node, [])
+        for i, kid in enumerate(kids):
+            render(kid, new_pref, i == len(kids) - 1, buf)
+        return buf
+
+    result = []
+    for r, root in enumerate(_roots(labels, children)):
+        result.append(render(root))
+    return result
+
+
+def build_task_table(result):
+    id2name, awaits = _index(result)
+    table = []
+    for tid, tasks in result:
+        for task_id, task_name, awaited in tasks:
+            for stack, awaiter_id in awaited:
+                coroutine_chain = " -> ".join(stack)
+                awaiter_name = id2name.get(awaiter_id, "Unknown")
+                table.append(
+                    [
+                        tid,
+                        hex(task_id),
+                        task_name,
+                        coroutine_chain,
+                        awaiter_name,
+                        hex(awaiter_id),
+                    ]
+                )
+
+    return table
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Show Python async tasks in a process")
+    parser.add_argument("pid", type=int, help="Process ID(s) to inspect.")
+    parser.add_argument(
+        "--tree", "-t", action="store_true", help="Display tasks in a tree format"
+    )
+    args = parser.parse_args()
+
+    try:
+        tasks = get_all_awaited_by(args.pid)
+    except RuntimeError as e:
+        print(f"Error retrieving tasks: {e}")
+        sys.exit(1)
+
+    if args.tree:
+        # Print the async call tree
+        result = print_async_tree(tasks)
+        for tree in result:
+            print("\n".join(tree))
+    else:
+        # Build and print the task table
+        table = build_task_table(tasks)
+        # Print the table in a simple tabular format
+        print(
+            f"{'tid':<10} {'task id':<20} {'task name':<20} {'coroutine chain':<50} {'awaiter name':<20} {'awaiter id':<15}"
+        )
+        print("-" * 135)
+        for row in table:
+            print(f"{row[0]:<10} {row[1]:<20} {row[2]:<20} {row[3]:<50} {row[4]:<20} {row[5]:<15}")
diff --git a/Lib/test/test_external_inspection.py b/Lib/test/test_external_inspection.py
@@ -13,13 +13,13 @@
 PROCESS_VM_READV_SUPPORTED = False
 
 try:
-    from _testexternalinspection import PROCESS_VM_READV_SUPPORTED
-    from _testexternalinspection import get_stack_trace
-    from _testexternalinspection import get_async_stack_trace
-    from _testexternalinspection import get_all_awaited_by
+    from _remotedebuggingg import PROCESS_VM_READV_SUPPORTED
+    from _remotedebuggingg import get_stack_trace
+    from _remotedebuggingg import get_async_stack_trace
+    from _remotedebuggingg import get_all_awaited_by
 except ImportError:
     raise unittest.SkipTest(
-        "Test only runs when _testexternalinspection is available")
+        "Test only runs when _remotedebuggingmodule is available")
 
 def _make_test_script(script_dir, script_basename, source):
     to_return = make_script(script_dir, script_basename, source)

diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py
@@ -1960,7 +1960,7 @@ def _supports_remote_attaching():
     PROCESS_VM_READV_SUPPORTED = False
 
     try:
-        from _testexternalinspection import PROCESS_VM_READV_SUPPORTED
+        from _remotedebuggingmodule import PROCESS_VM_READV_SUPPORTED
     except ImportError:
         pass
 

@@ -286,7 +286,7 @@ PYTHONPATH=$(COREPYTHONPATH)
 #_testcapi _testcapimodule.c
 #_testimportmultiple _testimportmultiple.c
 #_testmultiphase _testmultiphase.c
-#_testexternalinspection _testexternalinspection.c
+#_remotedebuggingmodule _remotedebuggingmodule.c
 #_testsinglephase _testsinglephase.c
 
 # ---

@@ -33,6 +33,7 @@
 # Modules that should always be present (POSIX and Windows):
 @MODULE_ARRAY_TRUE@array arraymodule.c
 @MODULE__ASYNCIO_TRUE@_asyncio _asynciomodule.c
+@MODULE__REMOTEDEBUGGING_TRUE@_remotedebugging _remotedebuggingmodule.c
 @MODULE__BISECT_TRUE@_bisect _bisectmodule.c
 @MODULE__CSV_TRUE@_csv _csv.c
 @MODULE__HEAPQ_TRUE@_heapq _heapqmodule.c
@@ -186,7 +187,6 @@
 @MODULE__TESTIMPORTMULTIPLE_TRUE@_testimportmultiple _testimportmultiple.c
 @MODULE__TESTMULTIPHASE_TRUE@_testmultiphase _testmultiphase.c
 @MODULE__TESTSINGLEPHASE_TRUE@_testsinglephase _testsinglephase.c
-@MODULE__TESTEXTERNALINSPECTION_TRUE@_testexternalinspection _testexternalinspection.c
 @MODULE__CTYPES_TEST_TRUE@_ctypes_test _ctypes/_ctypes_test.c
 
 # Limited API template modules; must be built as shared modules.