"""Figure I/O management utilities.
Provides functions for saving Matplotlib figures in various formats and
rendering them as SVG or other image formats in Jupyter environments.
"""
from __future__ import annotations
__all__ = ["save_and_show", "save_formats", "show"]
import contextlib
import gzip
import io
import os
import warnings
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from tempfile import NamedTemporaryFile
from types import ModuleType
from typing import Any, TypeAlias
from xml.dom import minidom
from xml.parsers.expat import ExpatError
import matplotlib.pyplot as plt
from matplotlib.backend_bases import get_registered_canvas_class
from matplotlib.figure import Figure
# Image extensions matplotlib's savefig knows how to write — if a caller
# passes a path that already ends with one of these, we strip it so the
# requested ``formats`` are appended cleanly instead of producing
# ``name.png.png`` / ``name.png.svg``.
_KNOWN_IMAGE_SUFFIXES = frozenset(
{
".png",
".pdf",
".svg",
".svgz",
".eps",
".ps",
".jpg",
".jpeg",
".tif",
".tiff",
".webp",
".raw",
".rgba",
}
)
def _normalize_image_stem(image_stem: str, *, warn: bool = True) -> str:
"""Strip a trailing image suffix from ``image_stem`` if present.
Callers occasionally pass ``"out/chart.png"`` instead of the
documented ``"out/chart"`` (often by reusing a path variable
that already carries an extension). Without normalization
``save_formats(..., formats=("png", "svg"))`` would emit
``chart.png.png`` and ``chart.png.svg`` — silent file-naming
bugs that pollute output directories. Normalize once at the
boundary and emit a :class:`UserWarning` so call sites can be
cleaned up over time. ``warn=False`` computes the same path silently
while preparing an output plan before any dm save effects.
"""
suffix = Path(image_stem).suffix.lower()
if suffix in _KNOWN_IMAGE_SUFFIXES:
normalized = image_stem[: -len(suffix)]
if warn:
warnings.warn(
(
f"save_formats: image_stem {image_stem!r} ends with image "
f"suffix {suffix!r}; stripping to {normalized!r}. "
"Pass the path *without* an extension to silence this."
),
UserWarning,
stacklevel=3,
)
return normalized
return image_stem
def _with_reproducible_metadata(
caller_metadata: dict[str, Any] | None, drop_key: str
) -> dict[str, Any]:
"""Return savefig ``metadata`` that drops a churning timestamp field.
Matplotlib's SVG / PDF backends embed a wall-clock timestamp
(``Date`` in SVG's ``<dc:date>``, ``CreationDate`` in PDF) that makes
every re-render differ even when the plotted data is identical.
Setting the field to ``None`` omits it entirely, so an unchanged
figure serialises to a byte-identical file. The caller always wins:
if they passed a ``metadata`` dict that already sets ``drop_key`` it is
returned unchanged; otherwise a copy with ``drop_key=None`` added is
returned, leaving their other keys intact.
"""
if caller_metadata is None:
return {drop_key: None}
if drop_key in caller_metadata:
return caller_metadata
merged = dict(caller_metadata)
merged[drop_key] = None
return merged
def _save_deterministic_svgz(
fig: Figure, out: str, bbox_inches: str | None, svg_kwargs: dict[str, Any]
) -> None:
"""Write SVGZ with a fixed gzip header timestamp."""
buf = io.BytesIO()
fig.savefig(buf, format="svg", bbox_inches=bbox_inches, **svg_kwargs)
with (
open(out, "wb") as raw,
gzip.GzipFile(fileobj=raw, mode="wb", mtime=0) as compressed,
):
compressed.write(buf.getvalue())
@dataclass(frozen=True)
class _SvgOutput:
"""An ``svg`` output (``compressed``: ``svgz``) saved with a pinned id salt and no date."""
path: str
compressed: bool
@dataclass(frozen=True)
class _PdfOutput:
"""A ``pdf`` output saved without its creation timestamp."""
path: str
@dataclass(frozen=True)
class _PlainOutput:
"""Any other output, passed to ``savefig`` as given (matplotlib picks the writer)."""
path: str
_SaveFormat: TypeAlias = _SvgOutput | _PdfOutput | _PlainOutput
def _parse_save_format(fmt: str, stem: str) -> _SaveFormat:
"""Parse one ``formats`` entry into the output it asks for at ``stem.<fmt>``.
A leading ``.`` is dropped. Only exact lowercase ``svg``/``svgz``/``pdf``
get the reproducible treatment; any other spelling is a plain output.
Support for the format is checked separately by :func:`_check_save_format`.
"""
ext = fmt.lstrip(".")
path = f"{stem}.{ext}"
match ext:
case "svg" | "svgz":
return _SvgOutput(path, compressed=ext == "svgz")
case "pdf":
return _PdfOutput(path)
case _:
return _PlainOutput(path)
def _check_save_format(
fig: Figure, output: _SaveFormat, kwargs: dict[str, Any]
) -> None:
"""Reject unsupported effective formats without constructing a canvas.
Match Matplotlib's current-canvas-first selection and explicit backend
class check. Registered fallbacks may install their print method during
construction, which remains on the real save path. SVGZ saves inner SVG;
arbitrary kwargs, rendering and filesystem failures are not preflighted.
"""
out = output.path
compressed = isinstance(output, _SvgOutput) and output.compressed
effective = "svg" if compressed else kwargs.get("format")
if effective is None:
effective = os.path.splitext(out)[1][1:]
if effective == "":
effective = fig.canvas.get_default_filetype()
effective = effective.lower()
backend = kwargs.get("backend")
if backend is not None:
from matplotlib.backends.registry import backend_registry
# Matplotlib documents this loader but does not annotate its signature.
load_backend_module: Callable[[str], ModuleType] = (
backend_registry.load_backend_module
)
canvas_class = load_backend_module(backend).FigureCanvas
if not hasattr(canvas_class, f"print_{effective}"):
raise ValueError(
f"The {backend!r} backend does not support {effective} output"
)
elif not hasattr(fig.canvas, f"print_{effective}"):
if get_registered_canvas_class(effective) is None:
raise ValueError(
"Format {!r} is not supported (supported formats: {})".format(
effective,
", ".join(sorted(fig.canvas.get_supported_filetypes())),
)
)
def _write_output(
fig: Figure,
output: _SaveFormat,
bbox_inches: str | None,
salt: str,
kwargs: dict[str, Any],
) -> None:
"""Save ``fig`` to one planned output.
SVG/SVGZ drop ``Date`` from the metadata and, unless ``svg.hashsalt`` is
already set, salt element ids with ``salt``; SVGZ gets a fixed gzip
timestamp. PDF drops ``CreationDate``. A caller ``metadata`` key always
wins. Other outputs use ``savefig`` unchanged. Write errors propagate.
"""
caller_metadata = kwargs.get("metadata")
match output:
case _SvgOutput(path=out, compressed=compressed):
svg_kwargs = dict(kwargs)
svg_kwargs["metadata"] = _with_reproducible_metadata(
caller_metadata, "Date"
)
# Only pin the hashsalt when the caller hasn't set one globally
# — a user who configured ``svg.hashsalt`` for their own
# reproducible build keeps it.
salted = (
plt.rc_context({"svg.hashsalt": salt})
if plt.rcParams["svg.hashsalt"] is None
else contextlib.nullcontext()
)
with salted:
if compressed:
_save_deterministic_svgz(fig, out, bbox_inches, svg_kwargs)
else:
fig.savefig(out, bbox_inches=bbox_inches, **svg_kwargs)
case _PdfOutput(path=out):
pdf_kwargs = dict(kwargs)
pdf_kwargs["metadata"] = _with_reproducible_metadata(
caller_metadata, "CreationDate"
)
fig.savefig(out, bbox_inches=bbox_inches, **pdf_kwargs)
case _PlainOutput(path=out):
fig.savefig(out, bbox_inches=bbox_inches, **kwargs)
[docs]
def show(image_path: str, size: int = 600, unit: str = "pt") -> None:
"""Load an SVG image and display it at the specified size in a browser or Jupyter.
Parameters
----------
image_path : str
Path to the SVG image to display.
size : int, optional
Desired output width. Default is 600.
unit : str, optional
Unit for the width ('pt', 'px', etc.). Default is 'pt'.
Raises
------
ImportError
If IPython is not installed. ``show`` renders inline in Jupyter
via IPython, which is an optional extra — install it with
``pip install "dartwork-mpl[notebook]"``.
"""
try:
from IPython.display import HTML, SVG, display
except ImportError as exc: # pragma: no cover - exercised via mock
raise ImportError(
"dm.show() needs IPython for inline Jupyter display, which is "
"an optional extra. Install it with "
"'pip install \"dartwork-mpl[notebook]\"' "
"(or 'uv add \"dartwork-mpl[notebook]\"')."
) from exc
def _display_svg_html(svg_data: str) -> None:
"""Wrap the IPython HTML-display call to centralise typing."""
display(HTML(svg_data)) # type: ignore[no-untyped-call]
# ``show`` renders SVG only. A non-SVG path (``.png`` / ``.pdf`` or an
# extensionless file) makes IPython's ``SVG`` raise an XML ExpatError
# — either while reading the file here or when we re-parse the payload
# below. Convert both to one clear message pointing at
# ``save_and_show`` (which dispatches by format) instead of surfacing
# a raw XML error. The payload is matplotlib's own SVG backend
# output, never user-supplied XML, so XXE/billion-laughs don't apply.
try:
svg_obj = SVG(data=image_path) # type: ignore[no-untyped-call]
dom = minidom.parseString(svg_obj.data) # trusted dartwork SVG only
except ExpatError as exc:
raise ValueError(
f"dm.show() displays SVG only, but {image_path!r} is not valid "
"SVG. Save with a '.svg' path, or use dm.save_and_show() which "
"handles PNG/JPG/PDF too."
) from exc
desired_width = size
doc_el = dom.documentElement
width_attr = doc_el.getAttribute("width") if doc_el else ""
height_attr = doc_el.getAttribute("height") if doc_el else ""
try:
width = float(width_attr.replace(unit, ""))
height = float(height_attr.replace(unit, ""))
except ValueError:
_display_svg_html(svg_obj.data)
return
if width <= 0:
_display_svg_html(svg_obj.data)
return
aspect_ratio = height / width
desired_height = int(desired_width * aspect_ratio)
# Replace width attribute.
for w_str in (str(width), str(int(width))):
old = f'width="{w_str}{unit}"'
if old in svg_obj.data:
svg_obj.data = svg_obj.data.replace(
old, f'width="{desired_width}{unit}"'
)
break
# Replace height attribute.
for h_str in (str(height), str(int(height))):
old = f'height="{h_str}{unit}"'
if old in svg_obj.data:
svg_obj.data = svg_obj.data.replace(
old, f'height="{desired_height}{unit}"'
)
break
_display_svg_html(svg_obj.data)
[docs]
def save_and_show(
fig: Figure,
image_path: str | None = None,
size: int = 600,
unit: str = "pt",
*,
adopt_orphan_tick_font: bool | None = None,
close_figure: bool = True,
**kwargs: Any,
) -> None:
"""Save a figure to disk, then display it in a Jupyter or web environment.
Parameters
----------
fig : matplotlib.figure.Figure
The Matplotlib figure to save and display.
image_path : str | None, optional
Path to save the image. If None, a system temporary file is used.
size : int, optional
Display width. Default is 600.
unit : str, optional
Unit for the size ('pt', 'px', etc.). Default is 'pt'.
adopt_orphan_tick_font : bool | None, optional
If ``True``, apply :func:`~dartwork_mpl.layout.adopt_axis_label_font`
before saving so unlabeled axes' tick labels take the axis-label
font, matching :func:`save_formats`. Mutates the figure (see that
function's Notes). Default is ``None`` — the value is read from
:data:`dartwork_mpl.config.adopt_orphan_tick_font` (itself
defaulting to ``False``). Pass ``True`` / ``False`` explicitly to
override per call.
close_figure : bool, optional
If ``True`` (default), the figure is closed via :func:`plt.close`
after saving — matching the historical behaviour, where the function
was intended for one-shot "render-then-display" use in notebooks.
Pass ``False`` to keep the figure open so you can keep editing it
(e.g. add an annotation and resave with :func:`save_formats`).
``save_formats`` itself never closes; this keyword brings parity.
**kwargs
Additional keyword arguments passed to ``savefig``.
"""
if adopt_orphan_tick_font is None:
from ..config import config
adopt_orphan_tick_font = config.adopt_orphan_tick_font
if adopt_orphan_tick_font:
from ..layout import adopt_axis_label_font
adopt_axis_label_font(fig)
if image_path is None:
with NamedTemporaryFile(suffix=".svg", delete=False) as tmp:
tmp_path = tmp.name
try:
fig.savefig(tmp_path, bbox_inches=None, **kwargs)
if close_figure:
plt.close(fig)
show(tmp_path, size=size, unit=unit)
finally:
Path(tmp_path).unlink(missing_ok=True)
else:
Path(image_path).parent.mkdir(parents=True, exist_ok=True)
fig.savefig(image_path, bbox_inches=None, **kwargs)
# matplotlib appends the default format when the path has no
# recognised extension (``"first"`` -> ``"first.png"``), so
# resolve the file actually written *before* closing the figure
# (its canvas is the authoritative list of supported formats).
saved_path = _resolve_saved_path(image_path, kwargs, fig)
if close_figure:
plt.close(fig)
_display_saved(saved_path, size=size, unit=unit)
def _resolve_saved_path(
image_path: str, kwargs: dict[str, Any], fig: Figure
) -> str:
"""Return the path ``savefig`` actually wrote for ``image_path``.
``savefig`` appends ``format`` / ``rcParams["savefig.format"]`` when
the path has no recognised extension.
"""
suffix = Path(image_path).suffix.lstrip(".").lower()
known = set(fig.canvas.get_supported_filetypes())
if suffix in known:
return image_path
fmt = str(kwargs.get("format") or plt.rcParams["savefig.format"])
return f"{image_path}.{fmt.lstrip('.')}"
def _display_saved(path: str, *, size: int, unit: str) -> None:
"""Display a saved figure inline, dispatching on its file format."""
suffix = Path(path).suffix.lower()
if suffix == ".svg":
show(path, size=size, unit=unit)
return
if suffix in (".png", ".jpg", ".jpeg", ".gif", ".webp"):
try:
from IPython.display import Image, display
except ImportError as exc: # pragma: no cover - mirror show()
raise ImportError(
"dm.save_and_show() needs IPython for inline display; "
"install it with 'pip install \"dartwork-mpl[notebook]\"'."
) from exc
display(Image(filename=path, width=size)) # type: ignore[no-untyped-call]
return
# Non-inlineable formats (PDF, EPS, …) were still saved; there is no
# inline preview, so report the path rather than crashing.
warnings.warn(
f"Saved {path!r}; inline preview is only available for SVG and "
"raster formats, so nothing is displayed for this format.",
stacklevel=2,
)