threadtree: ensure frame snapshot registry cleanup after thread exit

easypy/threadtree.py

@@ -97,6 +97,15 @@ def wrapper(*args, **kwargs):
         try:
             return target(*args, **kwargs)
         finally:
+            # remove the snapshot connecting parent → this thread
+            _FRAME_SNAPSHOTS_REGISTRY.pop((parent_thread.uuid, thread.uuid), None)
+
+            # if this thread can also be a parent for other children later,
+            # remove any snapshots where it is the parent
+            to_delete = [key for key in _FRAME_SNAPSHOTS_REGISTRY if key[0] == thread.uuid]
+            for key in to_delete:
+                _FRAME_SNAPSHOTS_REGISTRY.pop(key, None)
+
             IDENT_TO_UUID.pop(thread.ident)
 
     return _orig_start_new_thread(wrapper, *args, **kwargs)

tests/test_threadtree.py

@@ -1,6 +1,6 @@
 import threading
 
-from easypy.threadtree import walk_frames
+from easypy.threadtree import walk_frames, _FRAME_SNAPSHOTS_REGISTRY
 
 
 def collect_frames():
@@ -53,3 +53,4 @@ def func_a():
     ]
     actual = [frame.f_code.co_name for frame in frames]
     assert actual == expected
+    assert len(_FRAME_SNAPSHOTS_REGISTRY) == 0, "Frame snapshots registry should be empty after threads have finished"