ENH: add a renderer to save python-profile prof files

pyinstrument/__main__.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import codecs
 import fnmatch
 import glob
 import json
@@ -88,9 +87,9 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
         action="store",
         type="string",
         help=(
-            "how the report should be rendered. One of: 'text', 'html', 'json', 'speedscope' or python "
-            "import path to a renderer class. Defaults to the appropriate format for the extension "
-            "if OUTFILE is given, otherwise, defaults to 'text'."
+            "how the report should be rendered. One of: 'text', 'html', 'json', 'speedscope', 'pstats', "
+            "or python import path to a renderer class. Defaults to the appropriate format "
+            "for the extension if OUTFILE is given, otherwise, defaults to 'text'."
         ),
         default=None,
     )
@@ -263,7 +262,7 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
     # open the output file
 
     if options.outfile:
-        f = codecs.open(options.outfile, "w", "utf-8")
+        f = open(options.outfile, "w", encoding="utf-8", errors="surrogateescape")
         should_close_f_after_writing = True
     else:
         f = sys.stdout
@@ -277,6 +276,12 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
         parser.error(e.args[0])
         exit(1)
 
+    if renderer.output_is_binary and not options.outfile and file_is_a_tty(f):
+        parser.error(
+            "Can't write binary output to a terminal. Redirect to a file or use --outfile."
+        )
+        exit(1)
+
     # get the session - execute code or load from disk
 
     if options.load_prev:
@@ -324,13 +329,6 @@ def dash_m_callback(option: str, opt: str, value: str, parser: optparse.OptionPa
 
         session = profiler.stop()
 
-    if options.outfile:
-        f = codecs.open(options.outfile, "w", "utf-8")
-        should_close_f_after_writing = True
-    else:
-        f = sys.stdout
-        should_close_f_after_writing = False
-
     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)
@@ -458,15 +456,17 @@ def get_renderer_class(renderer: str) -> type[renderers.Renderer]:
         return renderers.SpeedscopeRenderer
     elif renderer == "session":
         return renderers.SessionRenderer
+    elif renderer == "pstats":
+        return renderers.PstatsRenderer
     else:
         try:
             return object_with_import_path(renderer)
         except (ValueError, ModuleNotFoundError, AttributeError) as err:
             # ValueError means we failed to import this object
             raise OptionsParseError(
                 f"Failed to find renderer with name {renderer!r}.\n"
-                "Options are text, html, json, speedscope or a Python import path to a Renderer\n"
-                "class.\n"
+                "Options are text, html, json, speedscope, pstats or a Python\n"
+                "import path to a Renderer class.\n"
                 "\n"
                 f"Underlying error: {err}\n"
             )
@@ -488,6 +488,8 @@ def guess_renderer_from_outfile(outfile: str) -> str | None:
         return "json"
     elif ext == ".pyisession":
         return "session"
+    elif ext == ".pstats":
+        return "pstats"
     else:
         return None
 

pyinstrument/frame.py

@@ -120,11 +120,13 @@ def total_self_time(self) -> float:
         The total amount of self time in this frame (including self time recorded by SelfTimeFrame
         children, and await time from AwaitTimeFrame children)
         """
-        self_time = 0
 
-        for child in self.children:
-            if child.identifier in (SELF_TIME_FRAME_IDENTIFIER, AWAIT_FRAME_IDENTIFIER):
-                self_time += child.time
+        # self time is time in this frame, minus time in children
+        self_time = self.time
+        real_children = [c for c in self.children if not c.is_synthetic]
+
+        for child in real_children:
+            self_time -= child.time
 
         return self_time
 

pyinstrument/renderers/__init__.py

@@ -2,5 +2,6 @@
 from pyinstrument.renderers.console import ConsoleRenderer
 from pyinstrument.renderers.html import HTMLRenderer
 from pyinstrument.renderers.jsonrenderer import JSONRenderer
+from pyinstrument.renderers.pstatsrenderer import PstatsRenderer
 from pyinstrument.renderers.session import SessionRenderer
 from pyinstrument.renderers.speedscope import SpeedscopeRenderer

pyinstrument/renderers/base.py

@@ -22,6 +22,11 @@ class Renderer:
     Renderer output file extension without dot prefix. The default value is `txt`
     """
 
+    output_is_binary: bool = False
+    """
+    Whether the output of this renderer is binary data. The default value is `False`.
+    """
+
     def __init__(self):
         pass
 

pyinstrument/renderers/pstatsrenderer.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import marshal
+from typing import Any, Dict, Tuple
+
+from pyinstrument import processors
+from pyinstrument.frame import Frame
+from pyinstrument.renderers.base import FrameRenderer, ProcessorList
+from pyinstrument.session import Session
+
+# pyright: strict
+
+FrameKey = Tuple[str, int, str]
+CallerValue = Tuple[float, int, float, float]
+FrameValue = Tuple[float, int, float, float, Dict[FrameKey, CallerValue]]
+StatsDict = Dict[FrameKey, FrameValue]
+
+
+class PstatsRenderer(FrameRenderer):
+    """
+    Outputs a marshaled dict, containing processed frames in pstat format,
+    suitable for processing by gprof2dot and snakeviz.
+    """
+
+    output_file_extension = "pstats"
+    output_is_binary = True
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+
+    def frame_key(self, frame: Frame) -> FrameKey:
+        return (frame.file_path or "", frame.line_no or 0, frame.function)
+
+    def render_frame(self, frame: Frame | None, stats: StatsDict) -> None:
+        if frame is None:
+            return
+
+        key = self.frame_key(frame)
+
+        if key not in stats:
+            # create a new entry
+            # being a statistical profiler, we don't know the exact call time or
+            # number of calls, they're stubbed out
+            call_time = -1
+            number_calls = -1
+            total_time = 0
+            cumulative_time = 0
+            callers: dict[FrameKey, CallerValue] = {}
+        else:
+            call_time, number_calls, total_time, cumulative_time, callers = stats[key]
+
+        # update the total time and cumulative time
+        total_time += frame.total_self_time
+        cumulative_time += frame.time
+
+        if frame.parent:
+            parent_key = self.frame_key(frame.parent)
+            if parent_key not in callers:
+                p_call_time = -1
+                p_number_calls = -1
+                p_total_time = 0
+                p_cumulative_time = 0
+            else:
+                p_call_time, p_number_calls, p_total_time, p_cumulative_time = callers[parent_key]
+
+            p_total_time += frame.total_self_time
+            p_cumulative_time += frame.time
+
+            callers[parent_key] = p_call_time, p_number_calls, p_total_time, p_cumulative_time
+
+        stats[key] = (call_time, number_calls, total_time, cumulative_time, callers)
+
+        for child in frame.children:
+            if not frame.is_synthetic:
+                self.render_frame(child, stats)
+
+    def render(self, session: Session):
+        frame = self.preprocess(session.root_frame())
+
+        stats = {}
+        self.render_frame(frame, stats)
+
+        # marshal.dumps returns bytes, so we need to decode it to a string
+        # using surrogateescape
+        return marshal.dumps(stats).decode(encoding="utf-8", errors="surrogateescape")
+
+    def default_processors(self) -> ProcessorList:
+        return [
+            processors.remove_importlib,
+            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,
+        ]

test/test_overflow.py

@@ -5,7 +5,13 @@
 import pytest
 
 from pyinstrument import Profiler
-from pyinstrument.renderers import ConsoleRenderer, HTMLRenderer, JSONRenderer, SpeedscopeRenderer
+from pyinstrument.renderers import (
+    ConsoleRenderer,
+    HTMLRenderer,
+    JSONRenderer,
+    PstatsRenderer,
+    SpeedscopeRenderer,
+)
 
 # Utilities
 
@@ -60,3 +66,7 @@ def test_json(deep_profiler_session):
 
 def test_speedscope(deep_profiler_session):
     SpeedscopeRenderer().render(deep_profiler_session)
+
+
+def test_prof(deep_profiler_session):
+    PstatsRenderer().render(deep_profiler_session)

test/test_pstats_renderer.py

@@ -0,0 +1,84 @@
+import time
+from pstats import Stats
+from test.fake_time_util import fake_time
+from typing import Any
+
+import pytest
+
+from pyinstrument import Profiler
+from pyinstrument.renderers import PstatsRenderer
+
+
+def a():
+    b()
+    c()
+
+
+def b():
+    d()
+
+
+def c():
+    d()
+
+
+def d():
+    e()
+
+
+def e():
+    time.sleep(1)
+
+
+@pytest.fixture(scope="module")
+def profiler_session():
+    with fake_time():
+        profiler = Profiler()
+        profiler.start()
+
+        a()
+
+        profiler.stop()
+        return profiler.last_session
+
+
+def test_pstats_renderer(profiler_session, tmp_path):
+    fname = tmp_path / "test.pstats"
+    pstats = PstatsRenderer().render(profiler_session)
+    with open(fname, "w", encoding="utf-8", errors="surrogateescape") as fid:
+        fid.write(pstats)
+    stats: Any = Stats(str(fname))
+    # Sanity check
+    assert stats.total_tt > 0
+    # The graph is
+    # a() -> b() -> d() -> e() -> time.sleep()
+    #    \-> c() /
+    # so make sure d has callers of b, c, and that the times make sense
+
+    # in stats,
+    #   keys are tuples (file_path, line, func)
+    #   values are tuples (calltime, numcalls, selftime, cumtime, callers)
+    # in callers,
+    #   keys are the same as in stats
+    #   values are the same as stats but without callers
+
+    # check the time of d
+    d_key = [k for k in stats.stats.keys() if k[2] == "d"][0]
+    d_val = stats.stats[d_key]
+    d_cumtime = d_val[3]
+    assert d_cumtime == pytest.approx(2)
+
+    # check d's callers times are split
+    b_key = [k for k in stats.stats.keys() if k[2] == "b"][0]
+    c_key = [k for k in stats.stats.keys() if k[2] == "c"][0]
+    d_callers = d_val[4]
+    b_cumtime = d_callers[b_key][3]
+    c_cumtime = d_callers[c_key][3]
+    assert b_cumtime == pytest.approx(1)
+    assert c_cumtime == pytest.approx(1)
+
+    # check the time of e
+    e_key = [k for k in stats.stats.keys() if k[2] == "e"][0]
+    e_val = stats.stats[e_key]
+    e_cumtime = e_val[3]
+    assert e_cumtime == pytest.approx(2)