Add flat mode to console renderer

pyinstrument/renderers/console.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import time
-from typing import Any
+from typing import Any, Dict, List, Tuple
 
 import pyinstrument
 from pyinstrument import processors
@@ -22,22 +24,25 @@ def __init__(
         self,
         unicode: bool = False,
         color: bool = False,
+        flat: bool = False,
         time: LiteralStr["seconds", "percent_of_total"] = "seconds",
         **kwargs: Any,
-    ):
+    ) -> None:
         """
         :param unicode: Use unicode, like box-drawing characters in the output.
         :param color: Enable color support, using ANSI color sequences.
         :param time: How to display the duration of each frame - ``'seconds'`` or ``'percent_of_total'``
+        :param flat: Display a flat profile instead of a call graph.
         """
         super().__init__(**kwargs)
 
         self.unicode = unicode
         self.color = color
+        self.flat = flat
         self.colors = self.colors_enabled if color else self.colors_disabled
         self.time = time
 
-    def render(self, session: Session):
+    def render(self, session: Session) -> str:
         result = self.render_preamble(session)
 
         frame = self.preprocess(session.root_frame())
@@ -48,13 +53,16 @@ def render(self, session: Session):
 
         self.root_frame = frame
 
-        result += self.render_frame(self.root_frame)
+        if self.flat:
+            result += self.render_frame_flat(self.root_frame)
+        else:
+            result += self.render_frame(self.root_frame)
         result += "\n"
 
         return result
 
     # pylint: disable=W1401
-    def render_preamble(self, session: Session):
+    def render_preamble(self, session: Session) -> str:
         lines = [
             r"",
             r"  _     ._   __/__   _ _  _  _ _/_  ",
@@ -82,28 +90,7 @@ def render_frame(self, frame: Frame, indent: str = "", child_indent: str = "") -
             or frame.total_self_time > 0.2 * self.root_frame.time
             or frame in frame.group.exit_frames
         ):
-            if self.time == "percent_of_total":
-                percent = self.frame_proportion_of_total_time(frame) * 100
-                time_str = self._ansi_color_for_time(frame) + f"{percent:.0f}%" + self.colors.end
-            else:
-                time_str = self._ansi_color_for_time(frame) + f"{frame.time:.3f}" + self.colors.end
-
-            name_color = self._ansi_color_for_name(frame)
-
-            class_name = frame.class_name
-            if class_name:
-                name = f"{class_name}.{frame.function}"
-            else:
-                name = frame.function
-
-            result = "{indent}{time_str} {name_color}{name}{c.end}  {c.faint}{code_position}{c.end}\n".format(
-                indent=indent,
-                time_str=time_str,
-                name_color=name_color,
-                name=name,
-                code_position=frame.code_position_short,
-                c=self.colors,
-            )
+            result = f"{indent}{self.frame_description(frame)}\n"
 
             if self.unicode:
                 indents = {"├": "├─ ", "│": "│  ", "└": "└─ ", " ": "   "}
@@ -137,11 +124,65 @@ def render_frame(self, frame: Frame, indent: str = "", child_indent: str = "") -
 
         return result
 
-    def frame_proportion_of_total_time(self, frame: Frame):
-        return frame.time / self.root_frame.time
+    def render_frame_flat(self, frame: Frame) -> str:
+        def walk(frame: Frame):
+            frame_id_to_time[frame.identifier] = (
+                frame_id_to_time.get(frame.identifier, 0) + frame.total_self_time
+            )
+
+            frame_id_to_frame[frame.identifier] = frame
+
+            for child in frame.children:
+                walk(child)
+
+        frame_id_to_time: Dict[str, float] = {}
+        frame_id_to_frame: Dict[str, Frame] = {}
+
+        walk(frame)
+
+        id_time_pairs: List[Tuple[str, float]] = sorted(
+            frame_id_to_time.items(), key=(lambda item: item[1]), reverse=True
+        )
+
+        # remove nodes that represent less than 0.1% of the total time
+        id_time_pairs = [pair for pair in id_time_pairs if pair[1] / self.root_frame.time > 0.001]
+
+        result = ""
+
+        for frame_id, self_time in id_time_pairs:
+            result += self.frame_description(frame_id_to_frame[frame_id], override_time=self_time)
+            result += "\n"
+
+        return result
+
+    def frame_description(self, frame: Frame, *, override_time: float | None = None) -> str:
+        time = override_time if override_time is not None else frame.time
+        time_color = self._ansi_color_for_time(time)
+
+        if self.time == "percent_of_total":
+            time_str = f"{self.frame_proportion_of_total_time(time) * 100:.1f}%"
+        else:
+            time_str = f"{time:.3f}"
+
+        value_str = f"{time_color}{time_str}{self.colors.end}"
+
+        class_name = frame.class_name
+        if class_name:
+            function_name = f"{class_name}.{frame.function}"
+        else:
+            function_name = frame.function
+        function_color = self._ansi_color_for_name(frame)
+        function_str = f"{function_color}{function_name}{self.colors.end}"
+
+        code_position_str = f"{self.colors.faint}{frame.code_position_short}{self.colors.end}"
+
+        return f"{value_str} {function_str}  {code_position_str}"
+
+    def frame_proportion_of_total_time(self, time: float) -> float:
+        return time / self.root_frame.time
 
-    def _ansi_color_for_time(self, frame: Frame):
-        proportion_of_total = self.frame_proportion_of_total_time(frame)
+    def _ansi_color_for_time(self, time: float) -> str:
+        proportion_of_total = self.frame_proportion_of_total_time(time)
 
         if proportion_of_total > 0.6:
             return self.colors.red
@@ -152,7 +193,7 @@ def _ansi_color_for_time(self, frame: Frame):
         else:
             return self.colors.bright_green + self.colors.faint
 
-    def _ansi_color_for_name(self, frame: Frame):
+    def _ansi_color_for_name(self, frame: Frame) -> str:
         if frame.is_application_code:
             return self.colors.bg_dark_blue_255 + self.colors.white_255
         else: