mgplot.finalise_plot

Functions to finalise and save plots to the file system.

  1"""Functions to finalise and save plots to the file system."""
  2
  3import re
  4import unicodedata
  5from collections.abc import Callable, Sequence
  6from pathlib import Path
  7from typing import Any, Final, NotRequired, Unpack
  8
  9import matplotlib.pyplot as plt
 10import numpy as np
 11from matplotlib.axes import Axes
 12from matplotlib.figure import Figure, SubFigure
 13from matplotlib.lines import Line2D
 14from matplotlib.patches import Rectangle
 15from matplotlib.transforms import blended_transform_factory
 16from pandas import Period, PeriodIndex
 17
 18from mgplot.annotation_utils import resolve_annotation_collisions
 19from mgplot.axis_utils import get_period_axes, refresh_period_labels, register_period_axes
 20from mgplot.keyword_checking import BaseKwargs, report_kwargs, validate_kwargs
 21from mgplot.settings import get_setting
 22
 23# --- constants
 24ME: Final[str] = "finalise_plot"
 25MAX_FILENAME_LENGTH: Final[int] = 150
 26DEFAULT_MARGIN: Final[float] = 0.02
 27TIGHT_LAYOUT_PAD: Final[float] = 1.1
 28FOOTNOTE_FONTSIZE: Final[int] = 8
 29FOOTNOTE_FONTSTYLE: Final[str] = "italic"
 30FOOTNOTE_COLOR: Final[str] = "#999999"
 31ZERO_LINE_WIDTH: Final[float] = 0.66
 32ZERO_LINE_COLOR: Final[str] = "#555555"
 33ZERO_AXIS_ADJUSTMENT: Final[float] = 0.02
 34DEFAULT_FILE_TITLE_NAME: Final[str] = "plot"
 35# --- annotated axvline text
 36VLINE_TEXT_FONTSIZE: Final[str] = "xx-small"
 37VLINE_TEXT_ROTATION: Final[int] = 90
 38VLINE_TEXT_OFFSET: Final[float] = 2.0  # points, to the right of the line
 39VLINE_TEXT_PAD: Final[float] = 0.01  # axes fraction, in from the top/bottom
 40VLINE_AUTO_BAND: Final[float] = 0.02  # fraction of the x-span sampled either side
 41
 42
 43class FinaliseKwargs(BaseKwargs):
 44    """Keyword arguments for the finalise_plot function."""
 45
 46    # --- value options
 47    suptitle: NotRequired[str | None]
 48    title: NotRequired[str | None]
 49    xlabel: NotRequired[str | None]
 50    ylabel: NotRequired[str | None]
 51    xlim: NotRequired[tuple[float | int | Period, float | int | Period] | None]
 52    ylim: NotRequired[tuple[float, float] | None]
 53    xticks: NotRequired[list[float | int | Period] | None]
 54    yticks: NotRequired[list[float] | None]
 55    xscale: NotRequired[str | None]
 56    yscale: NotRequired[str | None]
 57    # --- splat options
 58    legend: NotRequired[bool | dict[str, Any] | None]
 59    axhspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
 60    axvspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
 61    axhline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
 62    axvline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
 63    # --- options for annotations
 64    lfooter: NotRequired[str]
 65    rfooter: NotRequired[str]
 66    lheader: NotRequired[str]
 67    rheader: NotRequired[str]
 68    # --- file/save options
 69    pre_tag: NotRequired[str]
 70    tag: NotRequired[str]
 71    filename: NotRequired[str]
 72    chart_dir: NotRequired[str]
 73    file_type: NotRequired[str]
 74    dpi: NotRequired[int]
 75    figsize: NotRequired[tuple[float, float]]
 76    show: NotRequired[bool]
 77    # --- other options
 78    preserve_lims: NotRequired[bool]
 79    remove_legend: NotRequired[bool]
 80    zero_y: NotRequired[bool]
 81    y0: NotRequired[bool]
 82    x0: NotRequired[bool]
 83    axisbelow: NotRequired[bool]
 84    dont_save: NotRequired[bool]
 85    dont_close: NotRequired[bool]
 86    axes_only: NotRequired[bool]
 87
 88
 89VALUE_KWARGS = (
 90    "title",
 91    "xlabel",
 92    "ylabel",
 93    "xlim",
 94    "ylim",
 95    "xticks",
 96    "yticks",
 97    "xscale",
 98    "yscale",
 99)
100SPLAT_KWARGS = (
101    "axhspan",
102    "axvspan",
103    "axhline",
104    "axvline",
105    "legend",  # needs to be last in this tuple
106)
107HEADER_FOOTER_KWARGS = (
108    "lfooter",
109    "rfooter",
110    "lheader",
111    "rheader",
112)
113
114
115def sanitize_filename(filename: str, max_length: int = MAX_FILENAME_LENGTH) -> str:
116    """Convert a string to a safe filename.
117
118    Args:
119        filename: The string to convert to a filename
120        max_length: Maximum length for the filename
121
122    Returns:
123        A safe filename string
124
125    """
126    if not filename:
127        return "untitled"
128
129    # Normalize unicode characters (e.g., é -> e)
130    filename = unicodedata.normalize("NFKD", filename)
131
132    # Remove non-ASCII characters
133    filename = filename.encode("ascii", "ignore").decode("ascii")
134
135    # Convert to lowercase
136    filename = filename.lower()
137
138    # Replace spaces and other separators with hyphens
139    filename = re.sub(r"[\s\-_]+", "-", filename)
140
141    # Remove unsafe characters, keeping only alphanumeric and hyphens
142    filename = re.sub(r"[^a-z0-9\-]", "", filename)
143
144    # Remove leading/trailing hyphens and collapse multiple hyphens
145    filename = re.sub(r"^-+|-+$", "", filename)
146    filename = re.sub(r"-+", "-", filename)
147
148    # Truncate to max length
149    if len(filename) > max_length:
150        filename = filename[:max_length].rstrip("-")
151
152    # Ensure we have a valid filename
153    return filename or "untitled"
154
155
156def make_legend(axes: Axes, *, legend: None | bool | dict[str, Any]) -> None:
157    """Create a legend for the plot."""
158    if legend is None or legend is False:
159        return
160
161    if legend is True:  # use the global default settings
162        legend = get_setting("legend")
163
164    if isinstance(legend, dict):
165        axes.legend(**legend)
166        return
167
168    print(f"Warning: expected dict argument for legend, but got {type(legend)}.")
169
170
171def apply_value_kwargs(axes: Axes, value_kwargs_: Sequence[str], **kwargs: Unpack[FinaliseKwargs]) -> None:
172    """Set matplotlib elements by name using Axes.set().
173
174    Tricky: some plotting functions may set the xlabel or ylabel.
175    So ... we will set these if a setting is explicitly provided. If no
176    setting is provided, we will set to None if they are not already set.
177    If they have already been set, we will not change them.
178
179    """
180    # --- preliminary
181    function: dict[str, Callable[[], str]] = {
182        "xlabel": axes.get_xlabel,
183        "ylabel": axes.get_ylabel,
184        "title": axes.get_title,
185    }
186
187    def fail() -> str:
188        return ""
189
190    # --- loop over potential value settings
191    for setting in value_kwargs_:
192        value = _convert_period_value(axes, setting, kwargs.get(setting))
193        if setting in kwargs:
194            # deliberately set, so we will action
195            axes.set(**{setting: value})
196            continue
197        required_to_set = ("title", "xlabel", "ylabel")
198        if setting not in required_to_set:
199            # not set - and not required - so we can skip
200            continue
201
202        # we will set these 'required_to_set' ones
203        # provided they are not already set
204        already_set = function.get(setting, fail)()
205        if already_set and value is None:
206            continue
207
208        # if we get here, we will set the value (implicitly to None)
209        axes.set(**{setting: value})
210
211
212_SplatValue = bool | dict[str, Any] | Sequence[dict[str, Any]] | None
213
214# Keys in each splat-method's kwargs that are x-axis coordinates — when the
215# plot uses a PeriodIndex the axis is mapped to Period ordinals, so a Period
216# passed here must be converted to its ordinal for matplotlib.
217_PERIOD_COORD_KEYS: Final[dict[str, tuple[str, ...]]] = {
218    "axvline": ("x",),
219    "axvspan": ("xmin", "xmax"),
220}
221
222
223def _convert_period_coords(axes: Axes, method_name: str, item: dict[str, Any]) -> dict[str, Any]:
224    """Return a copy of item with any Period x-coordinates replaced by ordinals.
225
226    If the axes was period-mapped by mgplot, the Period's freq must match the
227    axes' stashed freq — otherwise the ordinals live in different spaces.
228    On an axes with no stash we trust the programmer and just take .ordinal.
229    """
230    keys = _PERIOD_COORD_KEYS.get(method_name)
231    if not keys:
232        return item
233    stash = get_period_axes(axes)
234    stashed_freq = stash[0] if stash is not None else None
235    converted = dict(item)
236    for key in keys:
237        val = converted.get(key)
238        if isinstance(val, Period):
239            if stashed_freq is not None and val.freqstr != stashed_freq:
240                raise ValueError(
241                    f"{method_name} Period freq {val.freqstr!r} does not match axes freq {stashed_freq!r}",
242                )
243            if stashed_freq is not None:
244                # Widen the stash so later label refresh covers this coordinate,
245                # even if it falls outside the plotted data's ordinal range.
246                register_period_axes(axes, PeriodIndex([val]))
247            converted[key] = val.ordinal
248    return converted
249
250
251# Value-kwargs whose entries are x-axis coordinates — like axvline/axvspan, a
252# Period passed here must be converted to its ordinal on a period-mapped axes.
253_PERIOD_X_VALUE_KWARGS: Final[tuple[str, ...]] = ("xlim", "xticks")
254
255
256def _convert_period_value(axes: Axes, setting: str, value: object) -> object:
257    """Return value with any Period x-coordinates replaced by ordinals.
258
259    xlim is a 2-tuple and xticks a list; either may contain Periods when the
260    caller describes the axis in calendar terms. On a period-mapped axes the
261    Period freq must match the axes' stashed freq (see register_period_axes).
262    Non-x settings, None, and non-Period entries pass through unchanged.
263    """
264    if setting not in _PERIOD_X_VALUE_KWARGS or not isinstance(value, (tuple, list)):
265        return value
266    stash = get_period_axes(axes)
267    stashed_freq = stash[0] if stash is not None else None
268    converted: list[object] = []
269    for val in value:
270        if isinstance(val, Period):
271            if stashed_freq is not None and val.freqstr != stashed_freq:
272                raise ValueError(
273                    f"{setting} Period freq {val.freqstr!r} does not match axes freq {stashed_freq!r}",
274                )
275            if stashed_freq is not None:
276                # Widen the stash so the later label refresh covers this coordinate.
277                register_period_axes(axes, PeriodIndex([val]))
278            converted.append(val.ordinal)
279        else:
280            converted.append(val)
281    return tuple(converted) if isinstance(value, tuple) else converted
282
283
284# --- annotated vertical lines
285# "text", "loc" and "text_kwargs" in an axvline dict describe its label rather
286# than the line, and are popped before the dict is splatted into ax.axvline().
287# loc is a space-separated string mixing a vertical word ("auto"/"top"/"bottom")
288# with a side word ("left"/"right"); either may be omitted. The side picks which
289# flank of the line the rotated label sits on, VLINE_TEXT_OFFSET points away.
290_VLINE_EDGES: Final[tuple[str, ...]] = ("top", "bottom")
291_VLINE_SIDES: Final[tuple[str, ...]] = ("left", "right")
292
293
294def _parse_vline_loc(loc: object) -> tuple[str, str]:
295    """Split a loc string into (edge, side), order-free and each optional.
296
297    edge is "auto", "top" or "bottom"; side is "left" or "right". An omitted
298    edge falls back to "auto" (choose by the data) and an omitted side to
299    "right" (today's placement). Raises on an unrecognised or doubled word --
300    each is a typo, not a choice.
301    """
302    if not isinstance(loc, str):
303        raise TypeError(f"axvline 'loc' must be a string, got {type(loc)}")
304    edge: str | None = None
305    side: str | None = None
306    for tok in loc.split():
307        if tok == "auto" or tok in _VLINE_EDGES:
308            if edge is not None:
309                raise ValueError(f"axvline 'loc' names two vertical positions, got {loc!r}")
310            edge = tok
311        elif tok in _VLINE_SIDES:
312            if side is not None:
313                raise ValueError(f"axvline 'loc' names two sides, got {loc!r}")
314            side = tok
315        else:
316            raise ValueError(f"axvline 'loc' word {tok!r} not recognised in {loc!r}")
317    return edge or "auto", side or "right"
318
319
320def _pop_vline_text(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]] | None:
321    """Remove the label keys from an axvline dict and return them, or None if unlabelled.
322
323    Raises on a bad loc, a non-dict text_kwargs, or label options given
324    without any text to place -- each of which is a typo, not a choice.
325    """
326    text = item.pop("text", None)
327    loc = item.pop("loc", "auto")
328    text_kwargs = item.pop("text_kwargs", {})
329
330    if text is None or not str(text).strip():
331        if loc != "auto" or text_kwargs:
332            raise ValueError("axvline 'loc'/'text_kwargs' given without any 'text' to place")
333        return None
334    _parse_vline_loc(loc)  # validate now; raises on a bad loc string
335    if not isinstance(text_kwargs, dict):
336        raise TypeError(f"axvline 'text_kwargs' must be a dict, got {type(text_kwargs)}")
337    return str(text), loc, text_kwargs
338
339
340def _data_y_extent(axes: Axes, x: float) -> tuple[float, float] | None:
341    """Return the (min, max) y of plotted data in a narrow x-band around x.
342
343    A rotated label occupies only a sliver of the x-axis, so a narrow band is
344    what the eye actually judges. Lines are matched by transform identity:
345    axhline/axvline use blended transforms, so this skips them and measures
346    only real data. Rectangles cover bar plots; axvspan/axhspan patches are
347    likewise blended and skipped. Returns None when nothing is measurable.
348    """
349    left, right = axes.get_xlim()
350    half = abs(right - left) * VLINE_AUTO_BAND
351    low, high = x - half, x + half
352    found: list[float] = []
353
354    for line in axes.get_lines():
355        if not isinstance(line, Line2D) or line.get_transform() is not axes.transData:
356            continue
357        xdata = np.asarray(line.get_xdata(), dtype=float)
358        ydata = np.asarray(line.get_ydata(), dtype=float)
359        if xdata.size != ydata.size or xdata.size == 0:
360            continue
361        wanted = (xdata >= low) & (xdata <= high) & ~np.isnan(ydata)
362        if wanted.any():
363            found.extend((float(ydata[wanted].min()), float(ydata[wanted].max())))
364
365    for patch in axes.patches:
366        if not isinstance(patch, Rectangle) or patch.get_data_transform() is not axes.transData:
367            continue
368        x0, width = patch.get_x(), patch.get_width()
369        y0, height = patch.get_y(), patch.get_height()
370        if min(x0, x0 + width) > high or max(x0, x0 + width) < low:
371            continue
372        found.extend((min(y0, y0 + height), max(y0, y0 + height)))
373
374    if not found:
375        return None
376    return min(found), max(found)
377
378
379def _auto_vline_loc(axes: Axes, x: float) -> str:
380    """Pick the end of the axes with more room between the data and the limit."""
381    extent = _data_y_extent(axes, x)
382    if extent is None:
383        return "top"  # nothing measurable (empty band, or an unrecognised artist)
384    data_low, data_high = extent
385    bottom_lim, top_lim = axes.get_ylim()
386    if top_lim >= bottom_lim:
387        top_gap, bottom_gap = top_lim - data_high, data_low - bottom_lim
388    else:  # inverted y-axis: values decrease going up the display
389        top_gap, bottom_gap = data_low - top_lim, bottom_lim - data_high
390    return "top" if top_gap >= bottom_gap else "bottom"
391
392
393def _annotate_vline(axes: Axes, item: dict[str, Any], line: Line2D, spec: tuple[str, str, dict]) -> None:
394    """Place a rotated text label just to one side of a vertical line.
395
396    The side (left/right of the line) comes from loc; the label is anchored
397    with x in data coordinates and y in axes coordinates, so it stays pinned
398    to the top/bottom of the plot regardless of any later change to the
399    y-limits.
400    """
401    text, loc, text_kwargs = spec
402    x = item.get("x", 0)  # matches the matplotlib default for axvline
403    edge, side = _parse_vline_loc(loc)
404    if edge == "auto":
405        edge = _auto_vline_loc(axes, float(x))
406    y, valign = (1.0 - VLINE_TEXT_PAD, "top") if edge == "top" else (VLINE_TEXT_PAD, "bottom")
407    x_off, halign = (VLINE_TEXT_OFFSET, "left") if side == "right" else (-VLINE_TEXT_OFFSET, "right")
408
409    options: dict[str, Any] = {
410        "rotation": VLINE_TEXT_ROTATION,
411        "fontsize": VLINE_TEXT_FONTSIZE,
412        "color": line.get_color(),
413        "ha": halign,
414        "va": valign,
415    }
416    options.update(text_kwargs)
417    axes.annotate(
418        text,
419        xy=(x, y),
420        xycoords=blended_transform_factory(axes.transData, axes.transAxes),
421        xytext=(x_off, 0),
422        textcoords="offset points",
423        **options,
424    )
425
426
427def _apply_splat(axes: Axes, method_name: str, value: _SplatValue) -> None:
428    """Apply a single splat kwarg, which may be a dict or sequence of dicts."""
429    if value is None or value is False:
430        return
431
432    if value is True:  # use the global default settings
433        value = get_setting(method_name)
434
435    # normalise to a list of dicts
436    if isinstance(value, dict):
437        value = [value]
438
439    if isinstance(value, Sequence):
440        method = getattr(axes, method_name)
441        for item in value:
442            if not isinstance(item, dict):
443                print(f"Warning: expected dict in {method_name} sequence, but got {type(item)}.")
444                continue
445            converted = _convert_period_coords(axes, method_name, item)
446            # _convert_period_coords always copies, so popping is safe here
447            spec = _pop_vline_text(converted) if method_name == "axvline" else None
448            artist = method(**converted)
449            if spec is not None and isinstance(artist, Line2D):
450                _annotate_vline(axes, converted, artist, spec)
451    else:
452        print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
453
454
455def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None:
456    """Set matplotlib elements dynamically using setting_name and splat."""
457    for method_name in settings:
458        if method_name not in kwargs:
459            continue
460
461        if method_name == "legend":
462            legend_value = kwargs.get(method_name)
463            if isinstance(legend_value, (bool, dict, type(None))):
464                make_legend(axes, legend=legend_value)
465            else:
466                print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.")
467            continue
468
469        value = kwargs.get(method_name)
470        if value is None or isinstance(value, (bool, dict, Sequence)):
471            _apply_splat(axes, method_name, value)
472        else:
473            print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
474
475
476def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
477    """Set figure size and apply chart annotations.
478
479    No-op when axes_only=True: the work here is all figure-level (resize,
480    corner text) and would stomp on other panels in a multi-axes figure.
481    """
482    if kwargs.get("axes_only"):
483        return
484    fig = axes.figure
485    fig_size = kwargs.get("figsize", get_setting("figsize"))
486    if not isinstance(fig, SubFigure):
487        fig.set_size_inches(*fig_size)
488
489    annotations = {
490        "rfooter": (0.99, 0.001, "right", "bottom"),
491        "lfooter": (0.01, 0.001, "left", "bottom"),
492        "rheader": (0.99, 0.999, "right", "top"),
493        "lheader": (0.01, 0.999, "left", "top"),
494    }
495
496    for annotation in HEADER_FOOTER_KWARGS:
497        if annotation in kwargs:
498            x_pos, y_pos, h_align, v_align = annotations[annotation]
499            fig.text(
500                x_pos,
501                y_pos,
502                str(kwargs.get(annotation, "")),
503                ha=h_align,
504                va=v_align,
505                fontsize=FOOTNOTE_FONTSIZE,
506                fontstyle=FOOTNOTE_FONTSTYLE,
507                color=FOOTNOTE_COLOR,
508            )
509
510
511def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
512    """Apply settings found in kwargs, after plotting the data."""
513    apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs)
514
515
516def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
517    """Apply settings found in kwargs."""
518
519    def check_kwargs(name: str) -> bool:
520        return name in kwargs and bool(kwargs.get(name))
521
522    apply_value_kwargs(axes, VALUE_KWARGS, **kwargs)
523    apply_annotations(axes, **kwargs)
524
525    if check_kwargs("zero_y"):
526        bottom, top = axes.get_ylim()
527        adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT
528        if bottom > -adj:
529            axes.set_ylim(bottom=-adj)
530        if top < adj:
531            axes.set_ylim(top=adj)
532
533    if check_kwargs("y0"):
534        low, high = axes.get_ylim()
535        if low < 0 < high:
536            axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
537
538    if check_kwargs("x0"):
539        low, high = axes.get_xlim()
540        if low < 0 < high:
541            axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
542
543    if check_kwargs("axisbelow"):
544        axes.set_axisbelow(True)
545
546
547def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
548    """Save the figure to file."""
549    saving = not kwargs.get("dont_save", False)  # save by default
550    if not saving:
551        return
552
553    try:
554        chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir")))
555
556        # Ensure directory exists
557        chart_dir.mkdir(parents=True, exist_ok=True)
558
559        suptitle = kwargs.get("suptitle", "")
560        title = kwargs.get("title", "")
561        pre_tag = kwargs.get("pre_tag", "")
562        tag = kwargs.get("tag", "")
563        name_override = kwargs.get("filename", "")
564        name_title = name_override or suptitle or title
565        file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME)
566        file_type = kwargs.get("file_type", get_setting("file_type")).lower()
567        dpi = kwargs.get("dpi", get_setting("dpi"))
568
569        # Construct filename components safely
570        filename_parts = []
571        if pre_tag:
572            filename_parts.append(sanitize_filename(pre_tag))
573        filename_parts.append(file_title)
574        if tag:
575            filename_parts.append(sanitize_filename(tag))
576
577        # Join filename parts and add extension
578        filename = "-".join(filter(None, filename_parts))
579        filepath = chart_dir / f"{filename}.{file_type}"
580
581        fig.savefig(filepath, dpi=dpi)
582
583    except (
584        OSError,
585        PermissionError,
586        FileNotFoundError,
587        ValueError,
588        RuntimeError,
589        TypeError,
590        UnicodeError,
591    ) as e:
592        print(f"Error: Could not save plot to file: {e}")
593
594
595# - public functions for finalise_plot()
596
597
598def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
599    """Finalise and save plots to the file system.
600
601    The filename for the saved plot is constructed from the global
602    chart_dir, the plot's title, any specified tag text, and the
603    file_type for the plot.
604
605    Args:
606        axes: Axes - matplotlib axes object - required
607        kwargs: FinaliseKwargs
608
609    """
610    # --- check the kwargs
611    report_kwargs(caller=ME, **kwargs)
612    validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs)
613
614    # --- sanity checks
615    if len(axes.get_children()) < 1:
616        print(f"Warning: {ME}() called with an empty axes, which was ignored.")
617        return
618
619    # --- remember axis-limits should we need to restore thems
620    xlim, ylim = axes.get_xlim(), axes.get_ylim()
621
622    # margins
623    axes.margins(DEFAULT_MARGIN)
624    axes.autoscale(tight=False)  # This is problematic ...
625
626    apply_kwargs(axes, **kwargs)
627
628    # tight layout and save the figure
629    fig = axes.figure
630    axes_only = kwargs.get("axes_only", False)
631    if not axes_only and (suptitle := kwargs.get("suptitle")):
632        fig.suptitle(suptitle)
633    if kwargs.get("preserve_lims"):
634        # restore the original limits of the axes
635        axes.set_xlim(xlim)
636        axes.set_ylim(ylim)
637    if not axes_only and not isinstance(fig, SubFigure):
638        fig.tight_layout(pad=TIGHT_LAYOUT_PAD)
639    apply_late_kwargs(axes, **kwargs)
640    # axvspan/axvline in late_kwargs may have widened xlim beyond what
641    # set_labels() last saw; regenerate ticks from the updated view.
642    refresh_period_labels(axes)
643    # de-collide end-of-line annotations now that the layout/limits are final
644    resolve_annotation_collisions(axes)
645    legend = axes.get_legend()
646    if legend and kwargs.get("remove_legend", False):
647        legend.remove()
648    if not axes_only and not isinstance(fig, SubFigure):
649        save_to_file(fig, **kwargs)
650
651    # show the plot in Jupyter Lab
652    if not axes_only and kwargs.get("show"):
653        plt.show()
654
655    # And close - the figure this axes belongs to, not pyplot's current figure
656    if not axes_only and not kwargs.get("dont_close", False):
657        root = fig
658        while isinstance(root, SubFigure):
659            root = root.figure
660        plt.close(root)
ME: Final[str] = 'finalise_plot'
MAX_FILENAME_LENGTH: Final[int] = 150
DEFAULT_MARGIN: Final[float] = 0.02
TIGHT_LAYOUT_PAD: Final[float] = 1.1
FOOTNOTE_FONTSIZE: Final[int] = 8
FOOTNOTE_FONTSTYLE: Final[str] = 'italic'
FOOTNOTE_COLOR: Final[str] = '#999999'
ZERO_LINE_WIDTH: Final[float] = 0.66
ZERO_LINE_COLOR: Final[str] = '#555555'
ZERO_AXIS_ADJUSTMENT: Final[float] = 0.02
DEFAULT_FILE_TITLE_NAME: Final[str] = 'plot'
VLINE_TEXT_FONTSIZE: Final[str] = 'xx-small'
VLINE_TEXT_ROTATION: Final[int] = 90
VLINE_TEXT_OFFSET: Final[float] = 2.0
VLINE_TEXT_PAD: Final[float] = 0.01
VLINE_AUTO_BAND: Final[float] = 0.02
class FinaliseKwargs(mgplot.keyword_checking.BaseKwargs):
44class FinaliseKwargs(BaseKwargs):
45    """Keyword arguments for the finalise_plot function."""
46
47    # --- value options
48    suptitle: NotRequired[str | None]
49    title: NotRequired[str | None]
50    xlabel: NotRequired[str | None]
51    ylabel: NotRequired[str | None]
52    xlim: NotRequired[tuple[float | int | Period, float | int | Period] | None]
53    ylim: NotRequired[tuple[float, float] | None]
54    xticks: NotRequired[list[float | int | Period] | None]
55    yticks: NotRequired[list[float] | None]
56    xscale: NotRequired[str | None]
57    yscale: NotRequired[str | None]
58    # --- splat options
59    legend: NotRequired[bool | dict[str, Any] | None]
60    axhspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
61    axvspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
62    axhline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
63    axvline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
64    # --- options for annotations
65    lfooter: NotRequired[str]
66    rfooter: NotRequired[str]
67    lheader: NotRequired[str]
68    rheader: NotRequired[str]
69    # --- file/save options
70    pre_tag: NotRequired[str]
71    tag: NotRequired[str]
72    filename: NotRequired[str]
73    chart_dir: NotRequired[str]
74    file_type: NotRequired[str]
75    dpi: NotRequired[int]
76    figsize: NotRequired[tuple[float, float]]
77    show: NotRequired[bool]
78    # --- other options
79    preserve_lims: NotRequired[bool]
80    remove_legend: NotRequired[bool]
81    zero_y: NotRequired[bool]
82    y0: NotRequired[bool]
83    x0: NotRequired[bool]
84    axisbelow: NotRequired[bool]
85    dont_save: NotRequired[bool]
86    dont_close: NotRequired[bool]
87    axes_only: NotRequired[bool]

Keyword arguments for the finalise_plot function.

suptitle: NotRequired[str | None]
title: NotRequired[str | None]
xlabel: NotRequired[str | None]
ylabel: NotRequired[str | None]
xlim: NotRequired[tuple[float | int | pandas.Period, float | int | pandas.Period] | None]
ylim: NotRequired[tuple[float, float] | None]
xticks: NotRequired[list[float | int | pandas.Period] | None]
yticks: NotRequired[list[float] | None]
xscale: NotRequired[str | None]
yscale: NotRequired[str | None]
legend: NotRequired[bool | dict[str, Any] | None]
axhspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
axvspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
axhline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
axvline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
lfooter: NotRequired[str]
rfooter: NotRequired[str]
lheader: NotRequired[str]
rheader: NotRequired[str]
pre_tag: NotRequired[str]
tag: NotRequired[str]
filename: NotRequired[str]
chart_dir: NotRequired[str]
file_type: NotRequired[str]
dpi: NotRequired[int]
figsize: NotRequired[tuple[float, float]]
show: NotRequired[bool]
preserve_lims: NotRequired[bool]
remove_legend: NotRequired[bool]
zero_y: NotRequired[bool]
y0: NotRequired[bool]
x0: NotRequired[bool]
axisbelow: NotRequired[bool]
dont_save: NotRequired[bool]
dont_close: NotRequired[bool]
axes_only: NotRequired[bool]
VALUE_KWARGS = ('title', 'xlabel', 'ylabel', 'xlim', 'ylim', 'xticks', 'yticks', 'xscale', 'yscale')
SPLAT_KWARGS = ('axhspan', 'axvspan', 'axhline', 'axvline', 'legend')
def sanitize_filename(filename: str, max_length: int = 150) -> str:
116def sanitize_filename(filename: str, max_length: int = MAX_FILENAME_LENGTH) -> str:
117    """Convert a string to a safe filename.
118
119    Args:
120        filename: The string to convert to a filename
121        max_length: Maximum length for the filename
122
123    Returns:
124        A safe filename string
125
126    """
127    if not filename:
128        return "untitled"
129
130    # Normalize unicode characters (e.g., é -> e)
131    filename = unicodedata.normalize("NFKD", filename)
132
133    # Remove non-ASCII characters
134    filename = filename.encode("ascii", "ignore").decode("ascii")
135
136    # Convert to lowercase
137    filename = filename.lower()
138
139    # Replace spaces and other separators with hyphens
140    filename = re.sub(r"[\s\-_]+", "-", filename)
141
142    # Remove unsafe characters, keeping only alphanumeric and hyphens
143    filename = re.sub(r"[^a-z0-9\-]", "", filename)
144
145    # Remove leading/trailing hyphens and collapse multiple hyphens
146    filename = re.sub(r"^-+|-+$", "", filename)
147    filename = re.sub(r"-+", "-", filename)
148
149    # Truncate to max length
150    if len(filename) > max_length:
151        filename = filename[:max_length].rstrip("-")
152
153    # Ensure we have a valid filename
154    return filename or "untitled"

Convert a string to a safe filename.

Args: filename: The string to convert to a filename max_length: Maximum length for the filename

Returns: A safe filename string

def make_legend( axes: matplotlib.axes._axes.Axes, *, legend: None | bool | dict[str, Any]) -> None:
157def make_legend(axes: Axes, *, legend: None | bool | dict[str, Any]) -> None:
158    """Create a legend for the plot."""
159    if legend is None or legend is False:
160        return
161
162    if legend is True:  # use the global default settings
163        legend = get_setting("legend")
164
165    if isinstance(legend, dict):
166        axes.legend(**legend)
167        return
168
169    print(f"Warning: expected dict argument for legend, but got {type(legend)}.")

Create a legend for the plot.

def apply_value_kwargs( axes: matplotlib.axes._axes.Axes, value_kwargs_: Sequence[str], **kwargs: Unpack[FinaliseKwargs]) -> None:
172def apply_value_kwargs(axes: Axes, value_kwargs_: Sequence[str], **kwargs: Unpack[FinaliseKwargs]) -> None:
173    """Set matplotlib elements by name using Axes.set().
174
175    Tricky: some plotting functions may set the xlabel or ylabel.
176    So ... we will set these if a setting is explicitly provided. If no
177    setting is provided, we will set to None if they are not already set.
178    If they have already been set, we will not change them.
179
180    """
181    # --- preliminary
182    function: dict[str, Callable[[], str]] = {
183        "xlabel": axes.get_xlabel,
184        "ylabel": axes.get_ylabel,
185        "title": axes.get_title,
186    }
187
188    def fail() -> str:
189        return ""
190
191    # --- loop over potential value settings
192    for setting in value_kwargs_:
193        value = _convert_period_value(axes, setting, kwargs.get(setting))
194        if setting in kwargs:
195            # deliberately set, so we will action
196            axes.set(**{setting: value})
197            continue
198        required_to_set = ("title", "xlabel", "ylabel")
199        if setting not in required_to_set:
200            # not set - and not required - so we can skip
201            continue
202
203        # we will set these 'required_to_set' ones
204        # provided they are not already set
205        already_set = function.get(setting, fail)()
206        if already_set and value is None:
207            continue
208
209        # if we get here, we will set the value (implicitly to None)
210        axes.set(**{setting: value})

Set matplotlib elements by name using Axes.set().

Tricky: some plotting functions may set the xlabel or ylabel. So ... we will set these if a setting is explicitly provided. If no setting is provided, we will set to None if they are not already set. If they have already been set, we will not change them.

def apply_splat_kwargs( axes: matplotlib.axes._axes.Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None:
456def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None:
457    """Set matplotlib elements dynamically using setting_name and splat."""
458    for method_name in settings:
459        if method_name not in kwargs:
460            continue
461
462        if method_name == "legend":
463            legend_value = kwargs.get(method_name)
464            if isinstance(legend_value, (bool, dict, type(None))):
465                make_legend(axes, legend=legend_value)
466            else:
467                print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.")
468            continue
469
470        value = kwargs.get(method_name)
471        if value is None or isinstance(value, (bool, dict, Sequence)):
472            _apply_splat(axes, method_name, value)
473        else:
474            print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")

Set matplotlib elements dynamically using setting_name and splat.

def apply_annotations( axes: matplotlib.axes._axes.Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
477def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
478    """Set figure size and apply chart annotations.
479
480    No-op when axes_only=True: the work here is all figure-level (resize,
481    corner text) and would stomp on other panels in a multi-axes figure.
482    """
483    if kwargs.get("axes_only"):
484        return
485    fig = axes.figure
486    fig_size = kwargs.get("figsize", get_setting("figsize"))
487    if not isinstance(fig, SubFigure):
488        fig.set_size_inches(*fig_size)
489
490    annotations = {
491        "rfooter": (0.99, 0.001, "right", "bottom"),
492        "lfooter": (0.01, 0.001, "left", "bottom"),
493        "rheader": (0.99, 0.999, "right", "top"),
494        "lheader": (0.01, 0.999, "left", "top"),
495    }
496
497    for annotation in HEADER_FOOTER_KWARGS:
498        if annotation in kwargs:
499            x_pos, y_pos, h_align, v_align = annotations[annotation]
500            fig.text(
501                x_pos,
502                y_pos,
503                str(kwargs.get(annotation, "")),
504                ha=h_align,
505                va=v_align,
506                fontsize=FOOTNOTE_FONTSIZE,
507                fontstyle=FOOTNOTE_FONTSTYLE,
508                color=FOOTNOTE_COLOR,
509            )

Set figure size and apply chart annotations.

No-op when axes_only=True: the work here is all figure-level (resize, corner text) and would stomp on other panels in a multi-axes figure.

def apply_late_kwargs( axes: matplotlib.axes._axes.Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
512def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
513    """Apply settings found in kwargs, after plotting the data."""
514    apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs)

Apply settings found in kwargs, after plotting the data.

def apply_kwargs( axes: matplotlib.axes._axes.Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
517def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
518    """Apply settings found in kwargs."""
519
520    def check_kwargs(name: str) -> bool:
521        return name in kwargs and bool(kwargs.get(name))
522
523    apply_value_kwargs(axes, VALUE_KWARGS, **kwargs)
524    apply_annotations(axes, **kwargs)
525
526    if check_kwargs("zero_y"):
527        bottom, top = axes.get_ylim()
528        adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT
529        if bottom > -adj:
530            axes.set_ylim(bottom=-adj)
531        if top < adj:
532            axes.set_ylim(top=adj)
533
534    if check_kwargs("y0"):
535        low, high = axes.get_ylim()
536        if low < 0 < high:
537            axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
538
539    if check_kwargs("x0"):
540        low, high = axes.get_xlim()
541        if low < 0 < high:
542            axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
543
544    if check_kwargs("axisbelow"):
545        axes.set_axisbelow(True)

Apply settings found in kwargs.

def save_to_file( fig: matplotlib.figure.Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
548def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
549    """Save the figure to file."""
550    saving = not kwargs.get("dont_save", False)  # save by default
551    if not saving:
552        return
553
554    try:
555        chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir")))
556
557        # Ensure directory exists
558        chart_dir.mkdir(parents=True, exist_ok=True)
559
560        suptitle = kwargs.get("suptitle", "")
561        title = kwargs.get("title", "")
562        pre_tag = kwargs.get("pre_tag", "")
563        tag = kwargs.get("tag", "")
564        name_override = kwargs.get("filename", "")
565        name_title = name_override or suptitle or title
566        file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME)
567        file_type = kwargs.get("file_type", get_setting("file_type")).lower()
568        dpi = kwargs.get("dpi", get_setting("dpi"))
569
570        # Construct filename components safely
571        filename_parts = []
572        if pre_tag:
573            filename_parts.append(sanitize_filename(pre_tag))
574        filename_parts.append(file_title)
575        if tag:
576            filename_parts.append(sanitize_filename(tag))
577
578        # Join filename parts and add extension
579        filename = "-".join(filter(None, filename_parts))
580        filepath = chart_dir / f"{filename}.{file_type}"
581
582        fig.savefig(filepath, dpi=dpi)
583
584    except (
585        OSError,
586        PermissionError,
587        FileNotFoundError,
588        ValueError,
589        RuntimeError,
590        TypeError,
591        UnicodeError,
592    ) as e:
593        print(f"Error: Could not save plot to file: {e}")

Save the figure to file.

def finalise_plot( axes: matplotlib.axes._axes.Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
599def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
600    """Finalise and save plots to the file system.
601
602    The filename for the saved plot is constructed from the global
603    chart_dir, the plot's title, any specified tag text, and the
604    file_type for the plot.
605
606    Args:
607        axes: Axes - matplotlib axes object - required
608        kwargs: FinaliseKwargs
609
610    """
611    # --- check the kwargs
612    report_kwargs(caller=ME, **kwargs)
613    validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs)
614
615    # --- sanity checks
616    if len(axes.get_children()) < 1:
617        print(f"Warning: {ME}() called with an empty axes, which was ignored.")
618        return
619
620    # --- remember axis-limits should we need to restore thems
621    xlim, ylim = axes.get_xlim(), axes.get_ylim()
622
623    # margins
624    axes.margins(DEFAULT_MARGIN)
625    axes.autoscale(tight=False)  # This is problematic ...
626
627    apply_kwargs(axes, **kwargs)
628
629    # tight layout and save the figure
630    fig = axes.figure
631    axes_only = kwargs.get("axes_only", False)
632    if not axes_only and (suptitle := kwargs.get("suptitle")):
633        fig.suptitle(suptitle)
634    if kwargs.get("preserve_lims"):
635        # restore the original limits of the axes
636        axes.set_xlim(xlim)
637        axes.set_ylim(ylim)
638    if not axes_only and not isinstance(fig, SubFigure):
639        fig.tight_layout(pad=TIGHT_LAYOUT_PAD)
640    apply_late_kwargs(axes, **kwargs)
641    # axvspan/axvline in late_kwargs may have widened xlim beyond what
642    # set_labels() last saw; regenerate ticks from the updated view.
643    refresh_period_labels(axes)
644    # de-collide end-of-line annotations now that the layout/limits are final
645    resolve_annotation_collisions(axes)
646    legend = axes.get_legend()
647    if legend and kwargs.get("remove_legend", False):
648        legend.remove()
649    if not axes_only and not isinstance(fig, SubFigure):
650        save_to_file(fig, **kwargs)
651
652    # show the plot in Jupyter Lab
653    if not axes_only and kwargs.get("show"):
654        plt.show()
655
656    # And close - the figure this axes belongs to, not pyplot's current figure
657    if not axes_only and not kwargs.get("dont_close", False):
658        root = fig
659        while isinstance(root, SubFigure):
660            root = root.figure
661        plt.close(root)

Finalise and save plots to the file system.

The filename for the saved plot is constructed from the global chart_dir, the plot's title, any specified tag text, and the file_type for the plot.

Args: axes: Axes - matplotlib axes object - required kwargs: FinaliseKwargs