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: