Make colour handling more consistent in the HTML instrumentation report

src/dvsim/instrumentation/report/base.py

@@ -9,6 +9,7 @@
 from pathlib import Path
 from typing import Any, Protocol, TypeVar
 
+import plotly.colors as pc
 import plotly.graph_objects as go
 import plotly.offline
 from typing_extensions import Self
@@ -24,6 +25,7 @@
     "PLOTLY_TIMING_AXIS_CONFIG",
     "InstrumentationVisualizer",
     "RenderProfile",
+    "get_default_color_map",
     "make_job_metadata_hover",
     "make_repeating_color_map",
     "render_html_report",
@@ -237,3 +239,22 @@ def make_repeating_color_map(data: Iterable[str], colors: Iterable[T]) -> dict[s
     while len(colors) < len(data):
         colors += colors.copy()
     return dict(zip(data, colors, strict=False))
+
+
+def get_default_color_map(categories: Sequence[str]) -> dict[str, Any]:
+    """Build a color map for a chart, using large palettes for more variety as is needed."""
+    palette = pc.qualitative.Plotly
+    if len(categories) > len(palette):
+        extra_colors = [
+            pc.qualitative.Bold,
+            pc.qualitative.Safe,
+            pc.qualitative.Vivid,
+            pc.qualitative.D3,
+            pc.qualitative.Set1,
+            pc.qualitative.Set2,
+        ]
+        extra_index = 0
+        while len(categories) > len(palette) and extra_index < len(extra_colors):
+            palette += extra_colors[extra_index]
+            extra_index += 1
+    return make_repeating_color_map(categories, palette)

src/dvsim/instrumentation/report/breakdown.py

@@ -6,9 +6,7 @@
 
 from collections import defaultdict
 from collections.abc import Callable
-from typing import Any
 
-import plotly.colors as pc
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
@@ -18,7 +16,7 @@
     DEFAULT_VISUALIZATION_HEIGHT_PX,
     PLOTLY_TIMING_AXIS_CONFIG,
     InstrumentationVisualizer,
-    make_repeating_color_map,
+    get_default_color_map,
     render_plotly_figure,
 )
 from dvsim.utils import format_time_as_hms as format_time
@@ -55,24 +53,6 @@ def __init__(
         # Default margin layout information
         self.margins = {"t": 80, "b": 40, "l": 50, "r": 20}
 
-    def _get_color_map(self, categories: dict[str, list[str]]) -> dict[str, Any]:
-        """Build a colour map for the chart, using large palettes for more variety as is needed."""
-        palette = pc.qualitative.Plotly
-        if len(categories) > len(palette):
-            extra_colors = [
-                pc.qualitative.Bold,
-                pc.qualitative.Safe,
-                pc.qualitative.Vivid,
-                pc.qualitative.D3,
-                pc.qualitative.Set1,
-                pc.qualitative.Set2,
-            ]
-            extra_index = 0
-            while len(categories) > len(palette) and extra_index < len(extra_colors):
-                palette += extra_colors[extra_index]
-                extra_index += 1
-        return make_repeating_color_map(categories, palette)
-
     def render(self, results: InstrumentationResults) -> str | None:
         """Render a breakdown (pie/bar chart) from the instrumentation results as a HTML fragment.
 
@@ -96,7 +76,7 @@ def render(self, results: InstrumentationResults) -> str | None:
             return None
 
         categories = dict(sorted(categories.items()))
-        color_map = self._get_color_map(categories)
+        color_map = get_default_color_map(list(categories.keys()))
 
         # Determine the chart dimensions for the different subplots
         pie_chart_height = min(self.MIN_PIE_HEIGHT_PX, DEFAULT_VISUALIZATION_HEIGHT_PX)

src/dvsim/instrumentation/report/longest.py

@@ -9,7 +9,6 @@
 from dataclasses import dataclass
 from typing import Any
 
-import plotly.colors as pc
 import plotly.graph_objects as go
 from typing_extensions import Self
 
@@ -24,8 +23,8 @@
     PLOTLY_TIMING_AXIS_CONFIG,
     InstrumentationVisualizer,
     RenderProfile,
+    get_default_color_map,
     make_job_metadata_hover,
-    make_repeating_color_map,
     render_plotly_figure,
 )
 from dvsim.job.status import JobStatus
@@ -380,8 +379,11 @@ def render(self, results: InstrumentationResults) -> str | None:
 
         if self.color_fn is not None:
             color_map = {key: self.color_fn(key) for key in categories}
+        elif self.category_fn is not None:
+            # We don't need a colour for `self.ALL_KEY`, so exclude it.
+            color_map = get_default_color_map(categories[1:])
         else:
-            color_map = make_repeating_color_map(categories, pc.qualitative.Plotly)
+            color_map = get_default_color_map(categories)
 
         # For each category, create a bar trace & menu (visibility) toggle button.
         traces: list[go.Bar] = []

src/dvsim/instrumentation/report/timelines.py

@@ -18,7 +18,6 @@
 
 import matplotlib.patches as mpatches
 import matplotlib.pyplot as plt
-import plotly.colors as pc
 import plotly.graph_objects as go
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure as MplFigure
@@ -32,8 +31,8 @@
     PLOTLY_TIMING_AXIS_CONFIG,
     InstrumentationVisualizer,
     RenderProfile,
+    get_default_color_map,
     make_job_metadata_hover,
-    make_repeating_color_map,
     render_plotly_figure,
 )
 from dvsim.logging import log
@@ -168,7 +167,7 @@ def _get_figure_contents(self, results: InstrumentationResults) -> FigureData |
             key = "all" if metadata is None else metadata.target
             categories[key].append(job_id)
         categories = dict(sorted(categories.items()))
-        color_map = make_repeating_color_map(categories, pc.qualitative.Plotly)
+        color_map = get_default_color_map(list(categories.keys()))
 
         # Get the information for all bars & traces to render in the figure
         figure_info: dict[str, list[BarInfo]] = defaultdict(list)

src/dvsim/instrumentation/report/usage.py

@@ -7,7 +7,6 @@
 from collections import defaultdict
 from collections.abc import Callable
 
-import plotly.colors as pc
 import plotly.graph_objects as go
 from plotly.graph_objs import Figure
 
@@ -17,7 +16,7 @@
     DEFAULT_VISUALIZATION_HEIGHT_PX,
     PLOTLY_TIMING_AXIS_CONFIG,
     InstrumentationVisualizer,
-    make_repeating_color_map,
+    get_default_color_map,
     render_plotly_figure,
 )
 
@@ -85,7 +84,7 @@ def _build(self, results: InstrumentationResults) -> Figure | None:
             else:
                 categories["Concurrent Jobs"].append(job_id)
         categories = dict(sorted(categories.items()))
-        color_map = make_repeating_color_map(categories, pc.qualitative.Plotly)
+        color_map = get_default_color_map(list(categories.keys()))
 
         # For each group, get concurrency events and plot them as a trace.
         fig = go.Figure()