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] = 3.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_VLINE_LOCATIONS: Final[tuple[str, ...]] = ("auto", "top", "bottom")
288
289
290def _pop_vline_text(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]] | None:
291    """Remove the label keys from an axvline dict and return them, or None if unlabelled.
292
293    Raises on a bad loc, a non-dict text_kwargs, or label options given
294    without any text to place -- each of which is a typo, not a choice.
295    """
296    text = item.pop("text", None)
297    loc = item.pop("loc", "auto")
298    text_kwargs = item.pop("text_kwargs", {})
299
300    if text is None or not str(text).strip():
301        if loc != "auto" or text_kwargs:
302            raise ValueError("axvline 'loc'/'text_kwargs' given without any 'text' to place")
303        return None
304    if loc not in _VLINE_LOCATIONS:
305        raise ValueError(f"axvline 'loc' must be one of {_VLINE_LOCATIONS}, got {loc!r}")
306    if not isinstance(text_kwargs, dict):
307        raise TypeError(f"axvline 'text_kwargs' must be a dict, got {type(text_kwargs)}")
308    return str(text), loc, text_kwargs
309
310
311def _data_y_extent(axes: Axes, x: float) -> tuple[float, float] | None:
312    """Return the (min, max) y of plotted data in a narrow x-band around x.
313
314    A rotated label occupies only a sliver of the x-axis, so a narrow band is
315    what the eye actually judges. Lines are matched by transform identity:
316    axhline/axvline use blended transforms, so this skips them and measures
317    only real data. Rectangles cover bar plots; axvspan/axhspan patches are
318    likewise blended and skipped. Returns None when nothing is measurable.
319    """
320    left, right = axes.get_xlim()
321    half = abs(right - left) * VLINE_AUTO_BAND
322    low, high = x - half, x + half
323    found: list[float] = []
324
325    for line in axes.get_lines():
326        if not isinstance(line, Line2D) or line.get_transform() is not axes.transData:
327            continue
328        xdata = np.asarray(line.get_xdata(), dtype=float)
329        ydata = np.asarray(line.get_ydata(), dtype=float)
330        if xdata.size != ydata.size or xdata.size == 0:
331            continue
332        wanted = (xdata >= low) & (xdata <= high) & ~np.isnan(ydata)
333        if wanted.any():
334            found.extend((float(ydata[wanted].min()), float(ydata[wanted].max())))
335
336    for patch in axes.patches:
337        if not isinstance(patch, Rectangle) or patch.get_data_transform() is not axes.transData:
338            continue
339        x0, width = patch.get_x(), patch.get_width()
340        y0, height = patch.get_y(), patch.get_height()
341        if min(x0, x0 + width) > high or max(x0, x0 + width) < low:
342            continue
343        found.extend((min(y0, y0 + height), max(y0, y0 + height)))
344
345    if not found:
346        return None
347    return min(found), max(found)
348
349
350def _auto_vline_loc(axes: Axes, x: float) -> str:
351    """Pick the end of the axes with more room between the data and the limit."""
352    extent = _data_y_extent(axes, x)
353    if extent is None:
354        return "top"  # nothing measurable (empty band, or an unrecognised artist)
355    data_low, data_high = extent
356    bottom_lim, top_lim = axes.get_ylim()
357    if top_lim >= bottom_lim:
358        top_gap, bottom_gap = top_lim - data_high, data_low - bottom_lim
359    else:  # inverted y-axis: values decrease going up the display
360        top_gap, bottom_gap = data_low - top_lim, bottom_lim - data_high
361    return "top" if top_gap >= bottom_gap else "bottom"
362
363
364def _annotate_vline(axes: Axes, item: dict[str, Any], line: Line2D, spec: tuple[str, str, dict]) -> None:
365    """Place a rotated text label just to the right of a vertical line.
366
367    The label is anchored with x in data coordinates and y in axes
368    coordinates, so it stays pinned to the top/bottom of the plot regardless
369    of any later change to the y-limits.
370    """
371    text, loc, text_kwargs = spec
372    x = item.get("x", 0)  # matches the matplotlib default for axvline
373    if loc == "auto":
374        loc = _auto_vline_loc(axes, float(x))
375    y, valign = (1.0 - VLINE_TEXT_PAD, "top") if loc == "top" else (VLINE_TEXT_PAD, "bottom")
376
377    options: dict[str, Any] = {
378        "rotation": VLINE_TEXT_ROTATION,
379        "fontsize": VLINE_TEXT_FONTSIZE,
380        "color": line.get_color(),
381        "ha": "left",
382        "va": valign,
383    }
384    options.update(text_kwargs)
385    axes.annotate(
386        text,
387        xy=(x, y),
388        xycoords=blended_transform_factory(axes.transData, axes.transAxes),
389        xytext=(VLINE_TEXT_OFFSET, 0),
390        textcoords="offset points",
391        **options,
392    )
393
394
395def _apply_splat(axes: Axes, method_name: str, value: _SplatValue) -> None:
396    """Apply a single splat kwarg, which may be a dict or sequence of dicts."""
397    if value is None or value is False:
398        return
399
400    if value is True:  # use the global default settings
401        value = get_setting(method_name)
402
403    # normalise to a list of dicts
404    if isinstance(value, dict):
405        value = [value]
406
407    if isinstance(value, Sequence):
408        method = getattr(axes, method_name)
409        for item in value:
410            if not isinstance(item, dict):
411                print(f"Warning: expected dict in {method_name} sequence, but got {type(item)}.")
412                continue
413            converted = _convert_period_coords(axes, method_name, item)
414            # _convert_period_coords always copies, so popping is safe here
415            spec = _pop_vline_text(converted) if method_name == "axvline" else None
416            artist = method(**converted)
417            if spec is not None and isinstance(artist, Line2D):
418                _annotate_vline(axes, converted, artist, spec)
419    else:
420        print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
421
422
423def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None:
424    """Set matplotlib elements dynamically using setting_name and splat."""
425    for method_name in settings:
426        if method_name not in kwargs:
427            continue
428
429        if method_name == "legend":
430            legend_value = kwargs.get(method_name)
431            if isinstance(legend_value, (bool, dict, type(None))):
432                make_legend(axes, legend=legend_value)
433            else:
434                print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.")
435            continue
436
437        value = kwargs.get(method_name)
438        if value is None or isinstance(value, (bool, dict, Sequence)):
439            _apply_splat(axes, method_name, value)
440        else:
441            print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
442
443
444def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
445    """Set figure size and apply chart annotations.
446
447    No-op when axes_only=True: the work here is all figure-level (resize,
448    corner text) and would stomp on other panels in a multi-axes figure.
449    """
450    if kwargs.get("axes_only"):
451        return
452    fig = axes.figure
453    fig_size = kwargs.get("figsize", get_setting("figsize"))
454    if not isinstance(fig, SubFigure):
455        fig.set_size_inches(*fig_size)
456
457    annotations = {
458        "rfooter": (0.99, 0.001, "right", "bottom"),
459        "lfooter": (0.01, 0.001, "left", "bottom"),
460        "rheader": (0.99, 0.999, "right", "top"),
461        "lheader": (0.01, 0.999, "left", "top"),
462    }
463
464    for annotation in HEADER_FOOTER_KWARGS:
465        if annotation in kwargs:
466            x_pos, y_pos, h_align, v_align = annotations[annotation]
467            fig.text(
468                x_pos,
469                y_pos,
470                str(kwargs.get(annotation, "")),
471                ha=h_align,
472                va=v_align,
473                fontsize=FOOTNOTE_FONTSIZE,
474                fontstyle=FOOTNOTE_FONTSTYLE,
475                color=FOOTNOTE_COLOR,
476            )
477
478
479def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
480    """Apply settings found in kwargs, after plotting the data."""
481    apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs)
482
483
484def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
485    """Apply settings found in kwargs."""
486
487    def check_kwargs(name: str) -> bool:
488        return name in kwargs and bool(kwargs.get(name))
489
490    apply_value_kwargs(axes, VALUE_KWARGS, **kwargs)
491    apply_annotations(axes, **kwargs)
492
493    if check_kwargs("zero_y"):
494        bottom, top = axes.get_ylim()
495        adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT
496        if bottom > -adj:
497            axes.set_ylim(bottom=-adj)
498        if top < adj:
499            axes.set_ylim(top=adj)
500
501    if check_kwargs("y0"):
502        low, high = axes.get_ylim()
503        if low < 0 < high:
504            axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
505
506    if check_kwargs("x0"):
507        low, high = axes.get_xlim()
508        if low < 0 < high:
509            axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
510
511    if check_kwargs("axisbelow"):
512        axes.set_axisbelow(True)
513
514
515def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
516    """Save the figure to file."""
517    saving = not kwargs.get("dont_save", False)  # save by default
518    if not saving:
519        return
520
521    try:
522        chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir")))
523
524        # Ensure directory exists
525        chart_dir.mkdir(parents=True, exist_ok=True)
526
527        suptitle = kwargs.get("suptitle", "")
528        title = kwargs.get("title", "")
529        pre_tag = kwargs.get("pre_tag", "")
530        tag = kwargs.get("tag", "")
531        name_override = kwargs.get("filename", "")
532        name_title = name_override or suptitle or title
533        file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME)
534        file_type = kwargs.get("file_type", get_setting("file_type")).lower()
535        dpi = kwargs.get("dpi", get_setting("dpi"))
536
537        # Construct filename components safely
538        filename_parts = []
539        if pre_tag:
540            filename_parts.append(sanitize_filename(pre_tag))
541        filename_parts.append(file_title)
542        if tag:
543            filename_parts.append(sanitize_filename(tag))
544
545        # Join filename parts and add extension
546        filename = "-".join(filter(None, filename_parts))
547        filepath = chart_dir / f"{filename}.{file_type}"
548
549        fig.savefig(filepath, dpi=dpi)
550
551    except (
552        OSError,
553        PermissionError,
554        FileNotFoundError,
555        ValueError,
556        RuntimeError,
557        TypeError,
558        UnicodeError,
559    ) as e:
560        print(f"Error: Could not save plot to file: {e}")
561
562
563# - public functions for finalise_plot()
564
565
566def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
567    """Finalise and save plots to the file system.
568
569    The filename for the saved plot is constructed from the global
570    chart_dir, the plot's title, any specified tag text, and the
571    file_type for the plot.
572
573    Args:
574        axes: Axes - matplotlib axes object - required
575        kwargs: FinaliseKwargs
576
577    """
578    # --- check the kwargs
579    report_kwargs(caller=ME, **kwargs)
580    validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs)
581
582    # --- sanity checks
583    if len(axes.get_children()) < 1:
584        print(f"Warning: {ME}() called with an empty axes, which was ignored.")
585        return
586
587    # --- remember axis-limits should we need to restore thems
588    xlim, ylim = axes.get_xlim(), axes.get_ylim()
589
590    # margins
591    axes.margins(DEFAULT_MARGIN)
592    axes.autoscale(tight=False)  # This is problematic ...
593
594    apply_kwargs(axes, **kwargs)
595
596    # tight layout and save the figure
597    fig = axes.figure
598    axes_only = kwargs.get("axes_only", False)
599    if not axes_only and (suptitle := kwargs.get("suptitle")):
600        fig.suptitle(suptitle)
601    if kwargs.get("preserve_lims"):
602        # restore the original limits of the axes
603        axes.set_xlim(xlim)
604        axes.set_ylim(ylim)
605    if not axes_only and not isinstance(fig, SubFigure):
606        fig.tight_layout(pad=TIGHT_LAYOUT_PAD)
607    apply_late_kwargs(axes, **kwargs)
608    # axvspan/axvline in late_kwargs may have widened xlim beyond what
609    # set_labels() last saw; regenerate ticks from the updated view.
610    refresh_period_labels(axes)
611    # de-collide end-of-line annotations now that the layout/limits are final
612    resolve_annotation_collisions(axes)
613    legend = axes.get_legend()
614    if legend and kwargs.get("remove_legend", False):
615        legend.remove()
616    if not axes_only and not isinstance(fig, SubFigure):
617        save_to_file(fig, **kwargs)
618
619    # show the plot in Jupyter Lab
620    if not axes_only and kwargs.get("show"):
621        plt.show()
622
623    # And close - the figure this axes belongs to, not pyplot's current figure
624    if not axes_only and not kwargs.get("dont_close", False):
625        root = fig
626        while isinstance(root, SubFigure):
627            root = root.figure
628        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] = 3.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:
424def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None:
425    """Set matplotlib elements dynamically using setting_name and splat."""
426    for method_name in settings:
427        if method_name not in kwargs:
428            continue
429
430        if method_name == "legend":
431            legend_value = kwargs.get(method_name)
432            if isinstance(legend_value, (bool, dict, type(None))):
433                make_legend(axes, legend=legend_value)
434            else:
435                print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.")
436            continue
437
438        value = kwargs.get(method_name)
439        if value is None or isinstance(value, (bool, dict, Sequence)):
440            _apply_splat(axes, method_name, value)
441        else:
442            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:
445def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
446    """Set figure size and apply chart annotations.
447
448    No-op when axes_only=True: the work here is all figure-level (resize,
449    corner text) and would stomp on other panels in a multi-axes figure.
450    """
451    if kwargs.get("axes_only"):
452        return
453    fig = axes.figure
454    fig_size = kwargs.get("figsize", get_setting("figsize"))
455    if not isinstance(fig, SubFigure):
456        fig.set_size_inches(*fig_size)
457
458    annotations = {
459        "rfooter": (0.99, 0.001, "right", "bottom"),
460        "lfooter": (0.01, 0.001, "left", "bottom"),
461        "rheader": (0.99, 0.999, "right", "top"),
462        "lheader": (0.01, 0.999, "left", "top"),
463    }
464
465    for annotation in HEADER_FOOTER_KWARGS:
466        if annotation in kwargs:
467            x_pos, y_pos, h_align, v_align = annotations[annotation]
468            fig.text(
469                x_pos,
470                y_pos,
471                str(kwargs.get(annotation, "")),
472                ha=h_align,
473                va=v_align,
474                fontsize=FOOTNOTE_FONTSIZE,
475                fontstyle=FOOTNOTE_FONTSTYLE,
476                color=FOOTNOTE_COLOR,
477            )

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:
480def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
481    """Apply settings found in kwargs, after plotting the data."""
482    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:
485def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
486    """Apply settings found in kwargs."""
487
488    def check_kwargs(name: str) -> bool:
489        return name in kwargs and bool(kwargs.get(name))
490
491    apply_value_kwargs(axes, VALUE_KWARGS, **kwargs)
492    apply_annotations(axes, **kwargs)
493
494    if check_kwargs("zero_y"):
495        bottom, top = axes.get_ylim()
496        adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT
497        if bottom > -adj:
498            axes.set_ylim(bottom=-adj)
499        if top < adj:
500            axes.set_ylim(top=adj)
501
502    if check_kwargs("y0"):
503        low, high = axes.get_ylim()
504        if low < 0 < high:
505            axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
506
507    if check_kwargs("x0"):
508        low, high = axes.get_xlim()
509        if low < 0 < high:
510            axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
511
512    if check_kwargs("axisbelow"):
513        axes.set_axisbelow(True)

Apply settings found in kwargs.

def save_to_file( fig: matplotlib.figure.Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
516def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
517    """Save the figure to file."""
518    saving = not kwargs.get("dont_save", False)  # save by default
519    if not saving:
520        return
521
522    try:
523        chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir")))
524
525        # Ensure directory exists
526        chart_dir.mkdir(parents=True, exist_ok=True)
527
528        suptitle = kwargs.get("suptitle", "")
529        title = kwargs.get("title", "")
530        pre_tag = kwargs.get("pre_tag", "")
531        tag = kwargs.get("tag", "")
532        name_override = kwargs.get("filename", "")
533        name_title = name_override or suptitle or title
534        file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME)
535        file_type = kwargs.get("file_type", get_setting("file_type")).lower()
536        dpi = kwargs.get("dpi", get_setting("dpi"))
537
538        # Construct filename components safely
539        filename_parts = []
540        if pre_tag:
541            filename_parts.append(sanitize_filename(pre_tag))
542        filename_parts.append(file_title)
543        if tag:
544            filename_parts.append(sanitize_filename(tag))
545
546        # Join filename parts and add extension
547        filename = "-".join(filter(None, filename_parts))
548        filepath = chart_dir / f"{filename}.{file_type}"
549
550        fig.savefig(filepath, dpi=dpi)
551
552    except (
553        OSError,
554        PermissionError,
555        FileNotFoundError,
556        ValueError,
557        RuntimeError,
558        TypeError,
559        UnicodeError,
560    ) as e:
561        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:
567def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
568    """Finalise and save plots to the file system.
569
570    The filename for the saved plot is constructed from the global
571    chart_dir, the plot's title, any specified tag text, and the
572    file_type for the plot.
573
574    Args:
575        axes: Axes - matplotlib axes object - required
576        kwargs: FinaliseKwargs
577
578    """
579    # --- check the kwargs
580    report_kwargs(caller=ME, **kwargs)
581    validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs)
582
583    # --- sanity checks
584    if len(axes.get_children()) < 1:
585        print(f"Warning: {ME}() called with an empty axes, which was ignored.")
586        return
587
588    # --- remember axis-limits should we need to restore thems
589    xlim, ylim = axes.get_xlim(), axes.get_ylim()
590
591    # margins
592    axes.margins(DEFAULT_MARGIN)
593    axes.autoscale(tight=False)  # This is problematic ...
594
595    apply_kwargs(axes, **kwargs)
596
597    # tight layout and save the figure
598    fig = axes.figure
599    axes_only = kwargs.get("axes_only", False)
600    if not axes_only and (suptitle := kwargs.get("suptitle")):
601        fig.suptitle(suptitle)
602    if kwargs.get("preserve_lims"):
603        # restore the original limits of the axes
604        axes.set_xlim(xlim)
605        axes.set_ylim(ylim)
606    if not axes_only and not isinstance(fig, SubFigure):
607        fig.tight_layout(pad=TIGHT_LAYOUT_PAD)
608    apply_late_kwargs(axes, **kwargs)
609    # axvspan/axvline in late_kwargs may have widened xlim beyond what
610    # set_labels() last saw; regenerate ticks from the updated view.
611    refresh_period_labels(axes)
612    # de-collide end-of-line annotations now that the layout/limits are final
613    resolve_annotation_collisions(axes)
614    legend = axes.get_legend()
615    if legend and kwargs.get("remove_legend", False):
616        legend.remove()
617    if not axes_only and not isinstance(fig, SubFigure):
618        save_to_file(fig, **kwargs)
619
620    # show the plot in Jupyter Lab
621    if not axes_only and kwargs.get("show"):
622        plt.show()
623
624    # And close - the figure this axes belongs to, not pyplot's current figure
625    if not axes_only and not kwargs.get("dont_close", False):
626        root = fig
627        while isinstance(root, SubFigure):
628            root = root.figure
629        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