Repository: joerick/pyinstrument

Issue:
Feature Request: Add support for speedscope
Dear Authors,
Thank you so much for this package! It's very useful. I was hoping you could add support for speedscope like py-spy does, which is a very detailed flame chart application (hosted on https://www.speedscope.app/)

Discussion:
goxberry: @joerick Thanks! And thank you for your review and for writing this profiler! Your design has made it straightforward to add this new feature. As you've suggested, I've added `SpeedscopeRenderer` to the tests in `test/test_overflow.py` and `test/test_profiler.py`. Because the functions `long_function_a` and `long_function_b` within `test/test_profiler.py` call the `sleep` function, calling both `long_function_a` and `long_function_b` suffices to generate an event timeline with nested function calls, so I elected to make the `test_speedscope_output` test similar to the other `test_*_output` tests for simplicity.

As a result of testing, I noticed something peculiar in `SpeedscopeRenderer` and wanted to get your insight about what might be happening. When I run `test_profiler.test_speedscope_output` in a debugger and:

* set a breakpoint at `test/test_profiler.py:149` and continue until I reach it
* step into `pyinstrument/renderers/speedscope.py`
* set another breakpoint at `pyinstrument/renderers/speedscopy.py:220` and continue until I reach it

I notice that the value of `session.duration` at that line is less than the value of `self._event_time`, and I'm not sure why `session.duration < self._event_time`; I would have expected `session.duration >= self._event_time`.

You should be able to reproduce this phenomenon yourself by running in a terminal:
```bash
> pytest --trace ${PATH_TO_PYINSTRUMENT_REPO}/test/test_profiler.py -k "test_speedscope_output"
(Pdb) b 149
(Pdb) c
(Pdb) s    # should step into speedscope.py.__init__() here
(Pdb) b 220
(Pdb) c
(Pdb) p session.duration     # on my machine, this command prints the value 0.00041103363037109375
(Pdb) p self._event_time      # on my machine, this command prints 0.75, as expected, within a tolerance of 0.3
```

where the first `>` denotes a shell prompt, and the `(Pdb)` denotes a `pdb` debugger prompt.

Do you have any insight as to why this behavior might occur?
joerick: Hi @goxberry! Yes, I can offer an explanation. The session.duration is computed by calling `time.time()` before and after profiling. But we are using 'fake time', which is defined here:

https://github.com/joerick/pyinstrument/blob/96dcb80ea480a93fabe9912a9d0089da3485f088/test/fake_time_util.py#L23-L32

This monkeypatches `time.sleep` and some internal machinery, but not `time.time`. So in the reality the session duration is correct, very little time has passed, but the test has a simulated reality, where time has passed!

If this is causing a problem, you could edit the `fake_time` function to also monkeypatch `time.time` (returning `fake_clock.get_time` should do the trick). Or we can just ignore the difference :)

goxberry: Thanks for the explanation! 

> If this is causing a problem, you could edit the fake_time function to also monkeypatch time.time (returning fake_clock.get_time should do the trick). Or we can just ignore the difference :)

In the interest of pragmatism, let's ignore the difference. I'll revert the change in `SpeedscopeRenderer` regarding `session.duration` vs `self._event_time` and remove the assertion checking the `endValue` field in the Speedscope output.

Patch:
diff --git a/docs/guide.md b/docs/guide.md
index 383343f..14ec2e4 100644
--- a/docs/guide.md
+++ b/docs/guide.md
@@ -35,8 +35,8 @@ Here are the options you can use:
                             save to <outfile>
       -r RENDERER, --renderer=RENDERER
                             how the report should be rendered. One of: 'text',
-                            'html', 'json', or python import path to a renderer
-                            class
+                            'html', 'json', 'speedscope', or python import path
+                            to a renderer class
       -t, --timeline        render as a timeline - preserve ordering and don't
                             condense repeated calls
       --hide=EXPR           glob-style pattern matching the file paths whose
diff --git a/docs/reference.md b/docs/reference.md
index 5ddc127..59dabd0 100644
--- a/docs/reference.md
+++ b/docs/reference.md
@@ -52,6 +52,8 @@ Therefore, rendering can be customised by changing the ``processors`` property.
 .. autoclass:: pyinstrument.renderers.HTMLRenderer
 
 .. autoclass:: pyinstrument.renderers.JSONRenderer
+
+.. autoclass:: pyinstrument.renderers.SpeedscopeRenderer
 ```
 
 ### Processors
diff --git a/pyinstrument/__main__.py b/pyinstrument/__main__.py
index 50ca5ca..a4f5d4c 100644
--- a/pyinstrument/__main__.py
+++ b/pyinstrument/__main__.py
@@ -1,3 +1,5 @@
+# coding: utf-8
+
 from __future__ import annotations
 
 import codecs
@@ -82,7 +84,7 @@ def main():
         action="store",
         type="string",
         help=(
-            "how the report should be rendered. One of: 'text', 'html', 'json', or python "
+            "how the report should be rendered. One of: 'text', 'html', 'json', 'speedscope' or python "
             "import path to a renderer class"
         ),
         default="text",
@@ -328,6 +330,8 @@ def get_renderer_class(renderer: str) -> Type[renderers.Renderer]:
         return renderers.HTMLRenderer
     elif renderer == "json":
         return renderers.JSONRenderer
+    elif renderer == "speedscope":
+        return renderers.SpeedscopeRenderer
     else:
         return object_with_import_path(renderer)
 
diff --git a/pyinstrument/renderers/__init__.py b/pyinstrument/renderers/__init__.py
index fe26221..c301c4a 100644
--- a/pyinstrument/renderers/__init__.py
+++ b/pyinstrument/renderers/__init__.py
@@ -2,3 +2,4 @@ from pyinstrument.renderers.base import Renderer
 from pyinstrument.renderers.console import ConsoleRenderer
 from pyinstrument.renderers.html import HTMLRenderer
 from pyinstrument.renderers.jsonrenderer import JSONRenderer
+from pyinstrument.renderers.speedscope import SpeedscopeRenderer
diff --git a/pyinstrument/renderers/speedscope.py b/pyinstrument/renderers/speedscope.py
new file mode 100644
index 0000000..a770feb
--- /dev/null
+++ b/pyinstrument/renderers/speedscope.py
@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, Union
+
+from pyinstrument import processors
+from pyinstrument.frame import BaseFrame
+from pyinstrument.renderers.base import ProcessorList, Renderer
+from pyinstrument.session import Session
+
+# pyright: strict
+
+
+@dataclass(frozen=True, eq=True)
+class SpeedscopeFrame:
+    """
+    Data class to store data needed for speedscope's concept of a
+    frame, hereafter referred to as a "speedscope frame", as opposed to
+    a "pyinstrument frame". This type must be hashable in order to use
+    it as a dictionary key; a dictionary will be used to track unique
+    speedscope frames.
+    """
+
+    name: str | None
+    file: str | None
+    line: int | None
+
+
+class SpeedscopeEventType(Enum):
+    """Enum representing the only two types of speedscope frame events"""
+
+    OPEN = "O"
+    CLOSE = "C"
+
+
+@dataclass
+class SpeedscopeEvent:
+    """
+    Data class to store speedscope's concept of an "event", which
+    corresponds to opening or closing stack frames as functions or
+    methods are entered or exited.
+    """
+
+    type: SpeedscopeEventType
+    at: float
+    frame: int
+
+
+@dataclass
+class SpeedscopeProfile:
+    """
+    Data class to store speedscope's concept of a "profile".
+    """
+
+    name: str
+    events: list[SpeedscopeEvent]
+    end_value: float
+    start_value: float = 0.0
+    type: str = "evented"
+    unit: str = "seconds"
+
+
+@dataclass
+class SpeedscopeFile:
+    """
+    Data class encoding fields in speedscope's JSON file schema
+    """
+
+    name: str
+    profiles: list[SpeedscopeProfile]
+    shared: dict[str, list[SpeedscopeFrame]]
+    schema: str = "https://www.speedscope.app/file-format-schema.json"
+    active_profile_index: None = None
+    exporter: str = "pyinstrument"
+
+
+SpeedscopeFrameDictType = Dict[str, Union[str, int, None]]
+SpeedscopeEventDictType = Dict[str, Union[SpeedscopeEventType, float, int]]
+
+
+class SpeedscopeEncoder(json.JSONEncoder):
+    """
+    Encoder class used by json.dumps to serialize the various
+    speedscope data classes.
+    """
+
+    def default(self, o: Any) -> Any:
+        if isinstance(o, SpeedscopeFile):
+            return {
+                "$schema": o.schema,
+                "name": o.name,
+                "activeProfileIndex": o.active_profile_index,
+                "exporter": o.exporter,
+                "profiles": o.profiles,
+                "shared": o.shared,
+            }
+        if isinstance(o, SpeedscopeProfile):
+            return {
+                "type": o.type,
+                "name": o.name,
+                "unit": o.unit,
+                "startValue": o.start_value,
+                "endValue": o.end_value,
+                "events": o.events,
+            }
+        if isinstance(o, (SpeedscopeFrame, SpeedscopeEvent)):
+            d: SpeedscopeFrameDictType | SpeedscopeEventDictType = o.__dict__
+            return d
+        if isinstance(o, SpeedscopeEventType):
+            return o.value
+        return json.JSONEncoder.default(self, o)
+
+
+class SpeedscopeRenderer(Renderer):
+    """
+    Outputs a tree of JSON conforming to the speedscope schema documented at
+
+    wiki: https://github.com/jlfwong/speedscope/wiki/Importing-from-custom-sources
+    schema: https://www.speedscope.app/file-format-schema.json
+    spec: https://github.com/jlfwong/speedscope/blob/main/src/lib/file-format-spec.ts
+    example: https://github.com/jlfwong/speedscope/blob/main/sample/profiles/speedscope/0.0.1/simple.speedscope.json
+
+    """
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+
+        # Member holding a running total of wall clock time needed to
+        # compute the times at which events occur
+        self._event_time: float = 0.0
+
+        # Map of speedscope frames to speedscope frame indices, needed
+        # to construct evented speedscope profiles; exploits LIFO
+        # property of popinfo method in Python 3.7+ dictionaries. This
+        # dictionary is used to build up the "shared" JSON array in
+        # speedscope's schema.
+        self._frame_to_index: dict[SpeedscopeFrame, int] = {}
+
+    def render_frame(self, frame: BaseFrame | None) -> list[SpeedscopeEvent]:
+        """
+        Builds up a list of speedscope events that are used to populate the
+        "events" array in speedscope-formatted JSON.
+
+        This method has two notable side effects:
+
+        * it populates the self._frame_to_index dictionary that matches
+          speedscope frames with their positions in the "shared" array of
+          speedscope output; this dictionary will be used to write this
+          "shared" array in the render method
+
+        * it accumulates a running total of time elapsed by
+          accumulating the self_time spent in each pyinstrument frame;
+          this running total is used by speedscope events to construct
+          a flame chart.
+        """
+
+        # if frame is None, recursion bottoms out; no event frames
+        # need to be added
+        if frame is None:
+            return []
+
+        # Otherwise, form a speedscope frame and add it to the frame
+        # to index map if the frame is not already a key in that map.
+        sframe = SpeedscopeFrame(frame.function, frame.file_path, frame.line_no)
+        if sframe not in self._frame_to_index:
+            self._frame_to_index[sframe] = len(self._frame_to_index)
+
+        # Get the frame index and add a speedscope event corresponding
+        # to opening a stack frame.
+        sframe_index = self._frame_to_index[sframe]
+        open_event = SpeedscopeEvent(SpeedscopeEventType.OPEN, self._event_time, sframe_index)
+        events_array: list[SpeedscopeEvent] = [open_event]
+
+        # Add stack frame open and close events for all child frames
+        # of this frame.
+        for child in frame.children:
+            events_array.extend(self.render_frame(child))
+
+        # Update event time for closing this stack frame.
+        #
+        # If number of frames approaches 1e16 * desired accuracy
+        # level, consider using Neumaier-Kahan summation; improves
+        # worst-case relative accuracy of sum from O(num_summands *
+        # eps) to (2 * eps + O(num_summands * eps * eps)), where eps
+        # is IEEE-754 double precision unit roundoff, approximately
+        # 1e-16. Average case relative accuracy expressions replace
+        # num_summands with sqrt(num_summands). However, Kahan
+        # summation quadruples operation count of sum, and Neumaier
+        # variant also adds a branch & swap for each summand. Pairwise
+        # summation isn't an option here because a running total is
+        # needed.
+        self._event_time += frame.self_time
+
+        # Add event closing this stack frame.
+        close_event = SpeedscopeEvent(SpeedscopeEventType.CLOSE, self._event_time, sframe_index)
+        events_array.append(close_event)
+
+        return events_array
+
+    def render(self, session: Session):
+        frame = self.preprocess(session.root_frame())
+
+        id_: str = time.strftime("%Y-%m-%dT%H-%M-%S", time.localtime(session.start_time))
+        name: str = "CPU profile for {} at {}".format(session.program, id_)
+
+        sprofile_list: list[SpeedscopeProfile] = [
+            SpeedscopeProfile(name, self.render_frame(frame), session.duration)
+        ]
+
+        # Exploits Python 3.7+ dictionary property of iterating over
+        # keys in insertion order to build the list of speedscope
+        # frames.
+        sframe_list: list[SpeedscopeFrame] = [sframe for sframe in iter(self._frame_to_index)]
+
+        shared_dict = {"frames": sframe_list}
+        speedscope_file = SpeedscopeFile(name, sprofile_list, shared_dict)
+
+        return "%s\n" % json.dumps(speedscope_file, cls=SpeedscopeEncoder)
+
+    def default_processors(self) -> ProcessorList:
+        """
+        Default Processors for speedscope renderer; note that
+        processors.aggregate_repeated_calls is removed because
+        speedscope is a timeline-based format.
+        """
+        return [
+            processors.remove_importlib,
+            processors.merge_consecutive_self_time,
+            processors.group_library_frames_processor,
+            processors.remove_unnecessary_self_time_nodes,
+            processors.remove_irrelevant_nodes,
+        ]


Test patch:
diff --git a/test/test_overflow.py b/test/test_overflow.py
index 7f94414..e30381e 100644
--- a/test/test_overflow.py
+++ b/test/test_overflow.py
@@ -5,7 +5,7 @@ import time
 import pytest
 
 from pyinstrument import Profiler
-from pyinstrument.renderers import ConsoleRenderer, HTMLRenderer, JSONRenderer
+from pyinstrument.renderers import ConsoleRenderer, HTMLRenderer, JSONRenderer, SpeedscopeRenderer
 
 # Utilities
 
@@ -56,3 +56,7 @@ def test_html(deep_profiler_session):
 
 def test_json(deep_profiler_session):
     JSONRenderer().render(deep_profiler_session)
+
+
+def test_speedscope(deep_profiler_session):
+    SpeedscopeRenderer().render(deep_profiler_session)
diff --git a/test/test_profiler.py b/test/test_profiler.py
index 12388bc..ae8bb4f 100644
--- a/test/test_profiler.py
+++ b/test/test_profiler.py
@@ -1,4 +1,5 @@
 import asyncio
+import dataclasses
 import json
 import time
 from functools import partial
@@ -10,6 +11,7 @@ import trio
 
 from pyinstrument import Profiler, renderers
 from pyinstrument.frame import BaseFrame, Frame
+from pyinstrument.renderers.speedscope import SpeedscopeEvent, SpeedscopeEventType, SpeedscopeFrame
 from pyinstrument.session import Session
 
 from .util import assert_never, busy_wait, flaky_in_ci
@@ -138,6 +140,130 @@ def test_json_output():
     assert len(root_frame["children"]) == 2
 
 
+def test_speedscope_output():
+    with fake_time():
+        with Profiler() as profiler:
+            long_function_a()
+            long_function_b()
+
+    output_data = profiler.output(renderers.SpeedscopeRenderer())
+
+    output = json.loads(output_data)
+
+    file_level_schema_fields = {
+        "$schema",
+        "name",
+        "exporter",
+        "activeProfileIndex",
+        "profiles",
+        "shared",
+    }
+    for file_field in file_level_schema_fields:
+        assert file_field in output
+
+    assert output["$schema"] == "https://www.speedscope.app/file-format-schema.json"
+    assert "pyinstrument" in output["exporter"]
+    assert output["activeProfileIndex"] is None
+    assert "CPU profile" in output["name"]
+
+    assert "frames" in output["shared"]
+    speedscope_frame_list = output["shared"]["frames"]
+
+    # Distinct functions called stores indices in key-value pairs because function
+    # index lookup needed for Speedscope event list tests. Were we not testing
+    # the event list, the distinct functions called could be stored as a tuple.
+    distinct_functions_called = {
+        "test_speedscope_output": 0,
+        "long_function_a": 1,
+        "sleep": 2,
+        "long_function_b": 3,
+    }
+    assert len(speedscope_frame_list) == len(distinct_functions_called)
+    speedscope_frame_fields = tuple([field.name for field in dataclasses.fields(SpeedscopeFrame)])
+    for (function_name, frame_index) in distinct_functions_called.items():
+        for frame_field in speedscope_frame_fields:
+            assert frame_field in speedscope_frame_list[frame_index]
+        assert speedscope_frame_list[frame_index]["name"] == function_name
+
+    speedscope_profile_list = output["profiles"]
+    assert len(speedscope_profile_list) == 1
+    speedscope_profile = speedscope_profile_list[0]
+    speedscope_profile_fields = ("type", "name", "unit", "startValue", "endValue", "events")
+    for profile_field in speedscope_profile_fields:
+        assert profile_field in speedscope_profile
+
+    # speedscope_profile["endValue"] is not tested because a fake_time mock
+    # timer is used to replace time.time, and this mock causes session.duration
+    # to differ from self._event_time just before exiting SpeedscopeRenderer.render
+    assert speedscope_profile["type"] == "evented"
+    assert speedscope_profile["unit"] == "seconds"
+    assert speedscope_profile["startValue"] == 0.0
+
+    output_event_tuple = (
+        SpeedscopeEvent(
+            SpeedscopeEventType.OPEN,
+            0.0,
+            distinct_functions_called["test_speedscope_output"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.OPEN,
+            0.0,
+            distinct_functions_called["long_function_a"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.OPEN,
+            0.0,
+            distinct_functions_called["sleep"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.CLOSE,
+            0.25,
+            distinct_functions_called["sleep"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.CLOSE,
+            0.25,
+            distinct_functions_called["long_function_a"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.OPEN,
+            0.25,
+            distinct_functions_called["long_function_b"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.OPEN,
+            0.25,
+            distinct_functions_called["sleep"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.CLOSE,
+            0.75,
+            distinct_functions_called["sleep"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.CLOSE,
+            0.75,
+            distinct_functions_called["long_function_b"],
+        ),
+        SpeedscopeEvent(
+            SpeedscopeEventType.CLOSE,
+            0.75,
+            distinct_functions_called["test_speedscope_output"],
+        ),
+    )
+
+    speedscope_event_list = speedscope_profile["events"]
+    assert len(speedscope_event_list) == len(output_event_tuple)
+    speedscope_event_fields = tuple([field.name for field in dataclasses.fields(SpeedscopeEvent)])
+    for (event_index, speedscope_event) in enumerate(speedscope_event_list):
+        for event_field in speedscope_event_fields:
+            assert event_field in speedscope_event
+
+        assert speedscope_event["type"] == output_event_tuple[event_index].type.value
+        assert speedscope_event["frame"] == output_event_tuple[event_index].frame
+        assert speedscope_event["at"] == pytest.approx(output_event_tuple[event_index].at)
+
+
 def test_empty_profile():
     with Profiler() as profiler:
         pass