Support tracking multiple tasks

rosys/analysis/tracking.py

@@ -1,4 +1,6 @@
+import asyncio
 import uuid
+from collections import defaultdict
 from functools import wraps
 from typing import Callable, ParamSpec, TypeVar
 
@@ -11,23 +13,28 @@
 class Track:
 
     def __init__(self) -> None:
-        self.stack: dict[int, str] = {}
+        self.stacks: defaultdict[int, dict[int, str]] = defaultdict(dict)
 
     def __call__(self, f: Callable[_P, _T]) -> Callable[_P, _T]:
         @wraps(f)
         async def wrap(*args, **kwargs):
+            task_id = id(asyncio.current_task())
+            stack = self.stacks[task_id]
             uid = uuid.uuid4().int
-            self.stack[uid] = f.__name__
+            stack[uid] = f.__name__
             try:
                 return await f(*args, **kwargs)
             finally:
-                del self.stack[uid]
+                del stack[uid]
+                if not stack:
+                    del self.stacks[task_id]
         return wrap
 
-    def ui(self) -> ui.label:
-        label = ui.label()
-        ui.timer(0.5, lambda: label.set_text(' → '.join(self.stack.values())))
-        return label
+    def ui(self) -> ui.markdown:
+        markdown = ui.markdown()
+        ui.timer(0.1, lambda: markdown.set_content('\n\n'.join(' → '.join(stack.values()).replace('_', r'\_')
+                                                               for stack in self.stacks.values())))
+        return markdown
 
 
 track = Track()