Skip to content

fix(report): type matplotlib handles #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bamdadd merged 1 commit into from
Jul 14, 2026
Merged

fix(report): type matplotlib handles #6

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions src/steerbench/report.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,11 @@
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Literal
from typing import TYPE_CHECKING, Literal

if TYPE_CHECKING:
from matplotlib.artist import Artist
from matplotlib.figure import Figure

# How to pick the sweet spot from the effect shift off baseline.
SweetSpotDirection = Literal["increase", "decrease", "abs"]
Expand Down Expand Up @@ -535,7 +539,7 @@ def plot_layer_sensitivity(data: ReportData) -> bytes:

lines_effect, labels_effect = ax_effect.get_legend_handles_labels()
# de-duplicate repeated trap labels
seen: dict[str, object] = {}
seen: dict[str, Artist] = {}
for line, label in zip(lines_effect, labels_effect, strict=True):
seen.setdefault(label, line)
ax_effect.legend(list(seen.values()), list(seen.keys()), loc="best", fontsize=8)
Expand All @@ -544,14 +548,14 @@ def plot_layer_sensitivity(data: ReportData) -> bytes:
return _figure_to_png(fig)


def _figure_to_png(fig: object) -> bytes:
def _figure_to_png(fig: Figure) -> bytes:
"""Serialise a matplotlib figure to PNG bytes and close it."""
from io import BytesIO

import matplotlib.pyplot as plt

buf = BytesIO()
fig.savefig(buf, format="png", dpi=120) # type: ignore[attr-defined]
fig.savefig(buf, format="png", dpi=120)
plt.close(fig)
return buf.getvalue()

Expand Down