Merge pull request #255 from joerick/fix-group-removal

Fix erroneous [group frames hidden] appearing in the output sometimes

examples/tbhide_demo.py

@@ -0,0 +1,25 @@
+import time
+
+
+def D():
+    time.sleep(0.7)
+
+
+def C():
+    __tracebackhide__ = True
+    time.sleep(0.1)
+    D()
+
+
+def B():
+    __tracebackhide__ = True
+    time.sleep(0.1)
+    C()
+
+
+def A():
+    time.sleep(0.1)
+    B()
+
+
+A()

pyinstrument/frame.py

@@ -367,14 +367,18 @@ def libraries(self) -> list[str]:
     def frames(self) -> Sequence[Frame]:
         return tuple(self._frames)
 
-    # pylint: disable=W0212
     def add_frame(self, frame: Frame):
         if frame.group:
-            frame.group._frames.remove(frame)
+            frame.group.remove_frame(frame)
 
         self._frames.append(frame)
         frame.group = self
 
+    def remove_frame(self, frame: Frame):
+        assert frame.group is self
+        self._frames.remove(frame)
+        frame.group = None
+
     @property
     def exit_frames(self):
         """

pyinstrument/frame_ops.py

@@ -97,6 +97,11 @@ def delete_frame_from_tree(
     parent.absorbed_time += frame.absorbed_time
 
     frame.remove_from_parent()
+    # in this call, recursive is true, even when replace_with is 'children'.
+    # When replace_with is 'self_time' or 'nothing', that's what we want. But
+    # when it's 'children', by now, the children have been removed and added
+    # to the parent, so recursive is irrelevant.
+    remove_frame_from_groups(frame, recursive=True)
 
 
 def combine_frames(frame: Frame, into: Frame):
@@ -119,3 +124,24 @@ def combine_frames(frame: Frame, into: Frame):
 
     into.add_children(frame.children)
     frame.remove_from_parent()
+    remove_frame_from_groups(frame, recursive=False)
+
+
+def remove_frame_from_groups(frame: Frame, recursive: bool):
+    """
+    Removes frame from any groups that it is a member of. Should be used when
+    removing a frame from a tree, so groups don't keep references to removed
+    frames.
+    """
+    if recursive and frame.children:
+        for child in frame.children:
+            remove_frame_from_groups(child, recursive=True)
+
+    if frame.group:
+        group = frame.group
+        group.remove_frame(frame)
+
+        if len(group.frames) == 1:
+            # a group with only one frame is meaningless, we'll remove it
+            # entirely.
+            group.remove_frame(group.frames[0])

pyinstrument/renderers/console.py

@@ -212,10 +212,10 @@ def default_processors(self) -> ProcessorList:
             processors.remove_tracebackhide,
             processors.merge_consecutive_self_time,
             processors.aggregate_repeated_calls,
-            processors.group_library_frames_processor,
             processors.remove_unnecessary_self_time_nodes,
             processors.remove_irrelevant_nodes,
             processors.remove_first_pyinstrument_frames_processor,
+            processors.group_library_frames_processor,
         ]
 
     class colors_enabled:

pyinstrument/renderers/html.py

@@ -96,8 +96,8 @@ def default_processors(self) -> ProcessorList:
             processors.remove_tracebackhide,
             processors.merge_consecutive_self_time,
             processors.aggregate_repeated_calls,
-            processors.group_library_frames_processor,
             processors.remove_unnecessary_self_time_nodes,
             processors.remove_irrelevant_nodes,
             processors.remove_first_pyinstrument_frames_processor,
+            processors.group_library_frames_processor,
         ]

pyinstrument/renderers/jsonrenderer.py

@@ -80,8 +80,8 @@ def default_processors(self) -> ProcessorList:
             processors.remove_tracebackhide,
             processors.merge_consecutive_self_time,
             processors.aggregate_repeated_calls,
-            processors.group_library_frames_processor,
             processors.remove_unnecessary_self_time_nodes,
             processors.remove_irrelevant_nodes,
             processors.remove_first_pyinstrument_frames_processor,
+            processors.group_library_frames_processor,
         ]

pyinstrument/renderers/pstatsrenderer.py

@@ -90,7 +90,6 @@ def default_processors(self) -> ProcessorList:
             processors.remove_tracebackhide,
             processors.merge_consecutive_self_time,
             processors.aggregate_repeated_calls,
-            processors.group_library_frames_processor,
             processors.remove_unnecessary_self_time_nodes,
             processors.remove_irrelevant_nodes,
             processors.remove_first_pyinstrument_frames_processor,

pyinstrument/renderers/speedscope.py

@@ -236,7 +236,6 @@ def default_processors(self) -> ProcessorList:
             processors.remove_importlib,
             processors.remove_tracebackhide,
             processors.merge_consecutive_self_time,
-            processors.group_library_frames_processor,
             processors.remove_unnecessary_self_time_nodes,
             processors.remove_irrelevant_nodes,
             processors.remove_first_pyinstrument_frames_processor,