Fix runpy trimming when used at the command line (joerick#215)

* Fix runpy trimming when used at the command line


* Add debug printing

* Fix some trim version variances

* Change logic to allow some self time in the runpy machinery.

* Remove unused imports

examples/busy_wait.py

@@ -0,0 +1,21 @@
+import time
+
+
+def function_1():
+    pass
+
+
+def function_2():
+    pass
+
+
+def main():
+    start_time = time.time()
+
+    while time.time() < start_time + 0.25:
+        function_1()
+        function_2()
+
+
+if __name__ == "__main__":
+    main()

pyinstrument/__main__.py

@@ -14,8 +14,6 @@
 
 import pyinstrument
 from pyinstrument import Profiler, renderers
-from pyinstrument.frame import Frame
-from pyinstrument.processors import ProcessorOptions
 from pyinstrument.session import Session
 from pyinstrument.util import (
     file_is_a_tty,
@@ -333,10 +331,6 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
         f = sys.stdout
         should_close_f_after_writing = False
 
-    if isinstance(renderer, renderers.FrameRenderer):
-        # remove this frame from the trace
-        renderer.processors.append(remove_first_pyinstrument_frame_processor)
-
     if isinstance(renderer, renderers.HTMLRenderer) and not options.outfile and file_is_a_tty(f):
         # don't write HTML to a TTY, open in browser instead
         output_filename = renderer.open_in_browser(session)
@@ -534,28 +528,6 @@ def save_report_to_temp_storage(session: Session):
     return path, identifier
 
 
-# pylint: disable=W0613
-def remove_first_pyinstrument_frame_processor(
-    frame: Frame | None, options: ProcessorOptions
-) -> Frame | None:
-    """
-    The first frame when using the command line is always the __main__ function. I want to remove
-    that from the output.
-    """
-    if frame is None:
-        return None
-
-    if frame.file_path is None:
-        return frame
-
-    if "pyinstrument" in frame.file_path and len(frame.children) == 1:
-        frame = frame.children[0]
-        frame.remove_from_parent()
-        return frame
-
-    return frame
-
-
 class CommandLineOptions:
     """
     A type that codifies the `options` variable.

pyinstrument/processors.py

@@ -213,3 +213,66 @@ def remove_irrelevant_nodes(
         remove_irrelevant_nodes(child, options=options, total_time=total_time)
 
     return frame
+
+
+# pylint: disable=W0613
+def remove_first_pyinstrument_frames_processor(
+    frame: Frame | None, options: ProcessorOptions
+) -> Frame | None:
+    """
+    The first few frames when using the command line are the __main__ of
+    pyinstrument, the eval, and the 'runpy' module. I want to remove that from
+    the output.
+    """
+    if frame is None:
+        return None
+
+    # the initial pyinstrument frame
+    def is_initial_pyinstrument_frame(frame: Frame):
+        return (
+            frame.file_path is not None
+            and re.match(r".*pyinstrument[/\\]__main__.py", frame.file_path)
+            and len(frame.children) > 0
+        )
+
+    def is_exec_frame(frame: Frame):
+        return (
+            frame.proportion_of_parent > 0.8
+            and frame.file_path is not None
+            and "<string>" in frame.file_path
+            and len(frame.children) > 0
+        )
+
+    def is_runpy_frame(frame: Frame):
+        return (
+            frame.proportion_of_parent > 0.8
+            and frame.file_path is not None
+            and (re.match(r".*runpy.py", frame.file_path) or "<frozen runpy>" in frame.file_path)
+            and len(frame.children) > 0
+        )
+
+    result = frame
+
+    if not is_initial_pyinstrument_frame(result):
+        return frame
+
+    result = result.children[0]
+
+    if not is_exec_frame(result):
+        return frame
+
+    result = result.children[0]
+
+    if not is_runpy_frame(result):
+        return frame
+
+    # at this point we know we've matched the first few frames of a command
+    # line invocation. We'll trim some runpy frames and return.
+
+    while is_runpy_frame(result):
+        result = result.children[0]
+
+    # remove this frame from the parent to make it the new root frame
+    result.remove_from_parent()
+
+    return result

pyinstrument/renderers/console.py

@@ -167,6 +167,7 @@ def default_processors(self) -> ProcessorList:
             processors.group_library_frames_processor,
             processors.remove_unnecessary_self_time_nodes,
             processors.remove_irrelevant_nodes,
+            processors.remove_first_pyinstrument_frames_processor,
         ]
 
     class colors_enabled:

pyinstrument/renderers/html.py

@@ -97,4 +97,5 @@ def default_processors(self) -> ProcessorList:
             processors.group_library_frames_processor,
             processors.remove_unnecessary_self_time_nodes,
             processors.remove_irrelevant_nodes,
+            processors.remove_first_pyinstrument_frames_processor,
         ]

pyinstrument/renderers/jsonrenderer.py

@@ -85,4 +85,5 @@ def default_processors(self) -> ProcessorList:
             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

@@ -238,4 +238,5 @@ def default_processors(self) -> ProcessorList:
             processors.group_library_frames_processor,
             processors.remove_unnecessary_self_time_nodes,
             processors.remove_irrelevant_nodes,
+            processors.remove_first_pyinstrument_frames_processor,
         ]

test/test_cmdline.py

@@ -1,4 +1,5 @@
 import os
+import re
 import subprocess
 import sys
 from pathlib import Path
@@ -195,3 +196,28 @@ def test_interval(self, pyinstrument_invocation, tmp_path: Path):
 
         assert "busy_wait" in str(output)
         assert "do_nothing" in str(output)
+
+    def test_invocation_machinery_is_trimmed(self, pyinstrument_invocation, tmp_path: Path):
+        busy_wait_py = tmp_path / "busy_wait.py"
+        busy_wait_py.write_text(BUSY_WAIT_SCRIPT)
+
+        output = subprocess.check_output(
+            [
+                *pyinstrument_invocation,
+                "--show-all",
+                str(busy_wait_py),
+            ],
+            universal_newlines=True,
+        )
+
+        print("Output:")
+        print(output)
+
+        first_profiling_line = re.search(r"^\d+(\.\d+)?\s+([^\s]+)\s+(.*)", output, re.MULTILINE)
+        assert first_profiling_line
+
+        function_name = first_profiling_line.group(2)
+        location = first_profiling_line.group(3)
+
+        assert function_name == "<module>"
+        assert "busy_wait.py" in location