Merge pull request #100 from mgmacias95/GH-91048-tasks

Gh 91048 tasks

Author: pablogsal

Lib/asyncio/tools.py

@@ -11,6 +11,23 @@ class NodeType(Enum):
     TASK = 2
 
 
+class CycleFoundException(Exception):
+    """Raised when there is a cycle when drawing the call tree."""
+
+    def __init__(self, cycles, id2name):
+        super().__init__()
+        self.cycles = cycles
+        self.id2name = id2name
+
+    def __str__(self):
+        for c in self.cycles:
+            names = " → ".join(self.id2name.get(tid, hex(tid)) for tid in c)
+            return (
+                "ERROR: await-graph contains cycles – cannot print a tree!\n"
+                f"cycle: {names}"
+            )
+
+
 # ─── indexing helpers ───────────────────────────────────────────
 def _index(result):
     id2name, awaits = {}, []
@@ -56,9 +73,9 @@ def _cor_node(parent_key, frame_name):
 
 
 def _roots(id2label, children):
-    roots = [n for n, lbl in id2label.items() if lbl == "Task-1"]
-    if roots:
-        return roots
+    # 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]
 
@@ -80,7 +97,7 @@ def _find_cycles(graph):
     empty list if the graph is acyclic.
     """
     WHITE, GREY, BLACK = 0, 1, 2
-    color = {n: WHITE for n in graph}
+    color = defaultdict(lambda: WHITE)
     path, cycles = [], []
 
     def dfs(v):
@@ -108,6 +125,10 @@ def print_async_tree(result, task_emoji="(T)", cor_emoji="", printer=print):
     prefixing tasks with *task_emoji* and coroutine frames with *cor_emoji*.
     """
     id2name, awaits = _index(result)
+    g = _task_graph(awaits)
+    cycles = _find_cycles(g)
+    if cycles:
+        raise CycleFoundException(cycles, id2name)
     labels, children = _build_tree(id2name, awaits)
 
     def pretty(node):
@@ -166,20 +187,9 @@ def build_task_table(result):
         print(f"Error retrieving tasks: {e}")
         sys.exit(1)
 
+    print(tasks)
     if args.tree:
         # Print the async call tree
-        id2name, awaits = _index(tasks)
-        g = _task_graph(awaits)
-        cycles = _find_cycles(g)
-
-        if cycles:
-            print("ERROR: await-graph contains cycles – cannot print a tree!\n")
-            for c in cycles:
-                # pretty-print task names instead of bare ids
-                names = " → ".join(id2name.get(tid, hex(tid)) for tid in c)
-                print(f"  cycle: {names}")
-            sys.exit(1)
-
         result = print_async_tree(tasks)
         for tree in result:
             print("\n".join(tree))