feat: add graph mode

The new graph mode allows switching to a live flame graph
visualisation of the current thread data.

.flake8

@@ -1,6 +1,6 @@
 [flake8]
 select = ANN,B,B9,C,D,E,F,W,I
-ignore = ANN101,ANN102,B950,D100,D104,D107,E203,E402,E501,F401,W503,W606
+ignore = ANN101,ANN102,ANN401,B950,D100,D104,D107,E203,E402,E501,F401,W503,W606
 max-line-length = 80
 docstring-convention = google
 application-import-names = austin_tui

README.md

@@ -175,6 +175,20 @@ profiling data is still being captured and processed in the background, so that
 when the view is resumed, the latest figures are shown.
 
 
+## Graph mode
+
+A live flame graph visualisation of the current thread statistics can be
+displayed by pressing <kbd>G</kbd>. This might help with identifying the largest
+frames at a glance.
+
+<p align="center">
+  <img src="art/austin-tui-flamegraph.gif"
+       style="box-shadow: #111 0px 0px 16px;"
+       alt="Austin TUI - Live flame graph" />
+</p>
+
+To toggle back to the top view, simply press <kbd>G</kbd> again.
+
 ## Save statistics
 
 Peeking at a running Python application is nice but in many cases you would want

austin_tui/adapters.py

@@ -32,6 +32,7 @@
 from austin_tui.model.system import Percentage
 from austin_tui.model.system import SystemModel
 from austin_tui.view import View
+from austin_tui.widgets.graph import FlameGraphData
 from austin_tui.widgets.markup import AttrString
 from austin_tui.widgets.table import TableData
 
@@ -379,3 +380,48 @@ def _add_frame_stats(
             _add_frame_stats(children[-1], "└─ ", "   ", 0, thread_stats.children)
 
         return frame_stats
+
+
+class FlameGraphAdapter(Adapter):
+    """Flame graph data adapter."""
+
+    def transform(self) -> dict:
+        """Transform according to the right model."""
+        austin = self._model.frozen_austin if self._model.frozen else self._model.austin
+        system = self._model.frozen_system if self._model.frozen else self._model.system
+        return self._transform(austin, system)  # type: ignore[arg-type]
+
+    def _transform(
+        self, austin: AustinModel, system: Union[SystemModel, FrozenSystemModel]
+    ) -> dict:
+        thread_key = austin.threads[austin.current_thread]
+        pid, _, thread = thread_key.partition(":")
+
+        thread = austin.stats.processes[int(pid)].threads[thread]
+
+        cs = {}  # type: ignore[var-annotated]
+        total = thread.total.value
+        total_pct = min(int(total / system.duration / 1e4), 100)
+        data: FlameGraphData = {
+            f"THREAD {thread.label} ⏲️  {fmt_time(total)} ({total_pct}%)": (total, cs)
+        }
+        levels = [(c, cs) for c in thread.children.values()]
+        while levels:
+            level, c = levels.pop(0)
+            k = f"{level.label.function} ({level.label.filename})"
+            if k in c:
+                v, cs = c[k]
+                c[k] = (v + level.total.value, cs)
+            else:
+                cs = {}
+                c[k] = (level.total.value, cs)
+            levels.extend(((c, cs) for c in level.children.values()))
+
+        return data
+
+    def update(self, data: FlameGraphData) -> bool:
+        """Update the table."""
+        (header,) = data
+        return self._view.flamegraph.set_data(data) | self._view.graph_header.set_text(
+            " FLAME GRAPH FOR " + header
+        )

austin_tui/controller.py

@@ -33,6 +33,7 @@
 from austin_tui.adapters import CpuAdapter
 from austin_tui.adapters import CurrentThreadAdapter
 from austin_tui.adapters import DurationAdapter
+from austin_tui.adapters import FlameGraphAdapter
 from austin_tui.adapters import MemoryAdapter
 from austin_tui.adapters import ThreadDataAdapter
 from austin_tui.adapters import ThreadFullDataAdapter
@@ -65,9 +66,11 @@ class AustinTUIController:
     thread_data = ThreadDataAdapter
     thread_full_data = ThreadFullDataAdapter
     command_line = CommandLineAdapter
+    flamegraph = FlameGraphAdapter
 
     def __init__(self) -> None:
         self._full_mode = False
+        self._graph = False
         self._scaler = None
         self._formatter = None
         self._last_timestamp = 0
@@ -93,10 +96,13 @@ def set_thread_data(self) -> None:
         if not self.model.austin.threads:
             return
 
-        if self._full_mode:
-            self.thread_full_data()  # type: ignore[call-arg]
+        if self._graph:
+            self.flamegraph()  # type: ignore[call-arg]
         else:
-            self.thread_data()  # type: ignore[call-arg]
+            if self._full_mode:
+                self.thread_full_data()  # type: ignore[call-arg]
+            else:
+                self.thread_data()  # type: ignore[call-arg]
 
         self._last_timestamp = self.model.austin.stats.timestamp
 
@@ -113,8 +119,24 @@ def set_thread(self) -> bool:
 
         return True
 
+    def _add_flamegraph_palette(self) -> None:
+        colors = [196, 202, 214, 124, 160, 166, 208]
+        palette = self.view.palette
+
+        for i, color in enumerate(colors):
+            palette.add_color(f"fg{i}", 15, color)
+            palette.add_color(f"fgf{i}", color)
+
+        self.view.flamegraph.set_palette(
+            (
+                [palette.get_color(f"fg{i}") for i in range(len(colors))],
+                [palette.get_color(f"fgf{i}") for i in range(len(colors))],
+            )
+        )
+
     def start(self) -> None:
         """Start event."""
+        self._add_flamegraph_palette()
         self.view.open()
         self.view.submit_task(self.update_loop())
 
@@ -153,7 +175,10 @@ async def update_loop(self) -> None:
         """The UI update loop."""
         while not self.view._stopped and self.view.is_open and self.view.root_widget:
             if self.update():
-                self.view.table.draw()
+                if self._graph:
+                    self.view.flamegraph.draw()
+                else:
+                    self.view.table.draw()
 
             self.view.root_widget.refresh()
 
@@ -180,21 +205,32 @@ def _change_thread(self, direction: ThreadNav) -> bool:
     async def on_next_thread(self) -> bool:
         """Handle next thread event."""
         if self._change_thread(ThreadNav.NEXT):
-            self.view.table.draw()
-            self.view.stats_view.refresh()
+            if self._graph:
+                self.view.flamegraph.draw()
+                self.view.flame_view.refresh()
+            else:
+                self.view.table.draw()
+                self.view.stats_view.refresh()
             return True
         return False
 
     async def on_previous_thread(self) -> bool:
         """Handle previous thread event."""
         if self._change_thread(ThreadNav.PREV):
-            self.view.table.draw()
-            self.view.stats_view.refresh()
+            if self._graph:
+                self.view.flamegraph.draw()
+                self.view.flame_view.refresh()
+            else:
+                self.view.table.draw()
+                self.view.stats_view.refresh()
             return True
         return False
 
     async def on_full_mode_toggled(self, _: Any = None) -> bool:
         """Toggle full mode."""
+        if self._graph:
+            return False
+
         self._full_mode = not self._full_mode
         self.set_thread_data()
 
@@ -269,3 +305,16 @@ async def on_threshold_down(self, _: Any = None) -> bool:
         th = self._change_threshold(-0.01) * 100.0
         self.view.threshold.set_text(f"{th:.0f}%")
         return True
+
+    async def on_graph_toggled(self, _: Any = None) -> bool:
+        """Toggle graph visualisation."""
+        self._graph = not self._graph
+
+        self.view.dataview_selector.select(self._graph)
+
+        if self._graph:
+            self.flamegraph()  # type: ignore[call-arg]
+        else:
+            self.set_thread_data()
+
+        return True

austin_tui/view/austin.py

@@ -72,48 +72,76 @@ async def on_quit(self) -> bool:
 
     async def on_full_mode_toggled(self) -> bool:
         """Handle Full Mode toggle."""
+        if self.graph_cmd.state:
+            return False
+
         self.full_mode_cmd.toggle()
         return True
 
+    async def on_graph_toggled(self) -> bool:
+        """Handle graph visualisation toggling."""
+        self.graph_cmd.toggle()
+        self.dataview_selector.refresh()
+        if self.graph_cmd.state:
+            self.full_mode_cmd.set_color("disabled")
+        else:
+            self.full_mode_cmd.toggle()
+            self.full_mode_cmd.toggle()
+        return True
+
     async def on_save(self, data: Any = None) -> bool:
         """Handle Save event."""
         self.notification.set_text("Saving collected statistics ...")
         return True
 
     async def on_table_up(self, data: Any = None) -> bool:
         """Handle Up Arrow on the table widget."""
-        self.stats_view.scroll_up()
-        self.stats_view.refresh()
+        view = self.flame_view if self.graph_cmd.state else self.stats_view
+
+        view.scroll_up()
+        view.refresh()
         return False
 
     async def on_table_down(self, data: Any = None) -> bool:
         """Handle Down Arrow on the table widget."""
-        self.stats_view.scroll_down()
-        self.stats_view.refresh()
+        view = self.flame_view if self.graph_cmd.state else self.stats_view
+
+        view.scroll_down()
+        view.refresh()
         return False
 
     async def on_table_pgup(self, data: Any = None) -> bool:
         """Handle Page Up on the table widget."""
-        self.stats_view.scroll_page_up()
-        self.stats_view.refresh()
+        view = self.flame_view if self.graph_cmd.state else self.stats_view
+
+        view.scroll_page_up()
+        view.refresh()
         return False
 
     async def on_table_pgdown(self, data: Any = None) -> bool:
         """Handle Page Down on the table widget."""
-        self.stats_view.scroll_page_down()
-        self.stats_view.refresh()
+        view = self.flame_view if self.graph_cmd.state else self.stats_view
+
+        view.scroll_page_down()
+        view.refresh()
         return False
 
     async def on_table_home(self, _: Any = None) -> bool:
         """Handle Home key on the table widget."""
-        self.stats_view.top()
-        self.stats_view.refresh()
+        view = self.flame_view if self.graph_cmd.state else self.stats_view
+
+        view.top()
+        view.refresh()
+
         return False
 
     async def on_table_end(self, _: Any = None) -> bool:
         """Handle End key on the table widget."""
-        self.stats_view.bottom()
-        self.stats_view.refresh()
+        view = self.flame_view if self.graph_cmd.state else self.stats_view
+
+        view.bottom()
+        view.refresh()
+
         return False
 
     async def on_play_pause(self, _: Any = None) -> bool: