mgplot.finalise_plot

  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
 10from matplotlib.axes import Axes
 11from matplotlib.figure import Figure, SubFigure
 12from pandas import Period, PeriodIndex
 13
 14from mgplot.annotation_utils import resolve_annotation_collisions
 15from mgplot.axis_utils import get_period_axes, refresh_period_labels, register_period_axes
 16from mgplot.keyword_checking import BaseKwargs, report_kwargs, validate_kwargs
 17from mgplot.settings import get_setting
 18
 19# --- constants
 20ME: Final[str] = "finalise_plot"
 21MAX_FILENAME_LENGTH: Final[int] = 150
 22DEFAULT_MARGIN: Final[float] = 0.02
 23TIGHT_LAYOUT_PAD: Final[float] = 1.1
 24FOOTNOTE_FONTSIZE: Final[int] = 8
 25FOOTNOTE_FONTSTYLE: Final[str] = "italic"
 26FOOTNOTE_COLOR: Final[str] = "#999999"
 27ZERO_LINE_WIDTH: Final[float] = 0.66
 28ZERO_LINE_COLOR: Final[str] = "#555555"
 29ZERO_AXIS_ADJUSTMENT: Final[float] = 0.02
 30DEFAULT_FILE_TITLE_NAME: Final[str] = "plot"
 31
 32
 33class FinaliseKwargs(BaseKwargs):
 34    """Keyword arguments for the finalise_plot function."""
 35
 36    # --- value options
 37    suptitle: NotRequired[str | None]
 38    title: NotRequired[str | None]
 39    xlabel: NotRequired[str | None]
 40    ylabel: NotRequired[str | None]
 41    xlim: NotRequired[tuple[float | int | Period, float | int | Period] | None]
 42    ylim: NotRequired[tuple[float, float] | None]
 43    xticks: NotRequired[list[float | int | Period] | None]
 44    yticks: NotRequired[list[float] | None]
 45    xscale: NotRequired[str | None]
 46    yscale: NotRequired[str | None]
 47    # --- splat options
 48    legend: NotRequired[bool | dict[str, Any] | None]
 49    axhspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
 50    axvspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
 51    axhline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
 52    axvline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
 53    # --- options for annotations
 54    lfooter: NotRequired[str]
 55    rfooter: NotRequired[str]
 56    lheader: NotRequired[str]
 57    rheader: NotRequired[str]
 58    # --- file/save options
 59    pre_tag: NotRequired[str]
 60    tag: NotRequired[str]
 61    filename: NotRequired[str]
 62    chart_dir: NotRequired[str]
 63    file_type: NotRequired[str]
 64    dpi: NotRequired[int]
 65    figsize: NotRequired[tuple[float, float]]
 66    show: NotRequired[bool]
 67    # --- other options
 68    preserve_lims: NotRequired[bool]
 69    remove_legend: NotRequired[bool]
 70    zero_y: NotRequired[bool]
 71    y0: NotRequired[bool]
 72    x0: NotRequired[bool]
 73    axisbelow: NotRequired[bool]
 74    dont_save: NotRequired[bool]
 75    dont_close: NotRequired[bool]
 76    axes_only: NotRequired[bool]
 77
 78
 79VALUE_KWARGS = (
 80    "title",
 81    "xlabel",
 82    "ylabel",
 83    "xlim",
 84    "ylim",
 85    "xticks",
 86    "yticks",
 87    "xscale",
 88    "yscale",
 89)
 90SPLAT_KWARGS = (
 91    "axhspan",
 92    "axvspan",
 93    "axhline",
 94    "axvline",
 95    "legend",  # needs to be last in this tuple
 96)
 97HEADER_FOOTER_KWARGS = (
 98    "lfooter",
 99    "rfooter",
100    "lheader",
101    "rheader",
102)
103
104
105def sanitize_filename(filename: str, max_length: int = MAX_FILENAME_LENGTH) -> str:
106    """Convert a string to a safe filename.
107
108    Args:
109        filename: The string to convert to a filename
110        max_length: Maximum length for the filename
111
112    Returns:
113        A safe filename string
114
115    """
116    if not filename:
117        return "untitled"
118
119    # Normalize unicode characters (e.g., é -> e)
120    filename = unicodedata.normalize("NFKD", filename)
121
122    # Remove non-ASCII characters
123    filename = filename.encode("ascii", "ignore").decode("ascii")
124
125    # Convert to lowercase
126    filename = filename.lower()
127
128    # Replace spaces and other separators with hyphens
129    filename = re.sub(r"[\s\-_]+", "-", filename)
130
131    # Remove unsafe characters, keeping only alphanumeric and hyphens
132    filename = re.sub(r"[^a-z0-9\-]", "", filename)
133
134    # Remove leading/trailing hyphens and collapse multiple hyphens
135    filename = re.sub(r"^-+|-+$", "", filename)
136    filename = re.sub(r"-+", "-", filename)
137
138    # Truncate to max length
139    if len(filename) > max_length:
140        filename = filename[:max_length].rstrip("-")
141
142    # Ensure we have a valid filename
143    return filename or "untitled"
144
145
146def make_legend(axes: Axes, *, legend: None | bool | dict[str, Any]) -> None:
147    """Create a legend for the plot."""
148    if legend is None or legend is False:
149        return
150
151    if legend is True:  # use the global default settings
152        legend = get_setting("legend")
153
154    if isinstance(legend, dict):
155        axes.legend(**legend)
156        return
157
158    print(f"Warning: expected dict argument for legend, but got {type(legend)}.")
159
160
161def apply_value_kwargs(axes: Axes, value_kwargs_: Sequence[str], **kwargs: Unpack[FinaliseKwargs]) -> None:
162    """Set matplotlib elements by name using Axes.set().
163
164    Tricky: some plotting functions may set the xlabel or ylabel.
165    So ... we will set these if a setting is explicitly provided. If no
166    setting is provided, we will set to None if they are not already set.
167    If they have already been set, we will not change them.
168
169    """
170    # --- preliminary
171    function: dict[str, Callable[[], str]] = {
172        "xlabel": axes.get_xlabel,
173        "ylabel": axes.get_ylabel,
174        "title": axes.get_title,
175    }
176
177    def fail() -> str:
178        return ""
179
180    # --- loop over potential value settings
181    for setting in value_kwargs_:
182        value = _convert_period_value(axes, setting, kwargs.get(setting))
183        if setting in kwargs:
184            # deliberately set, so we will action
185            axes.set(**{setting: value})
186            continue
187        required_to_set = ("title", "xlabel", "ylabel")
188        if setting not in required_to_set:
189            # not set - and not required - so we can skip
190            continue
191
192        # we will set these 'required_to_set' ones
193        # provided they are not already set
194        already_set = function.get(setting, fail)()
195        if already_set and value is None:
196            continue
197
198        # if we get here, we will set the value (implicitly to None)
199        axes.set(**{setting: value})
200
201
202_SplatValue = bool | dict[str, Any] | Sequence[dict[str, Any]] | None
203
204# Keys in each splat-method's kwargs that are x-axis coordinates — when the
205# plot uses a PeriodIndex the axis is mapped to Period ordinals, so a Period
206# passed here must be converted to its ordinal for matplotlib.
207_PERIOD_COORD_KEYS: Final[dict[str, tuple[str, ...]]] = {
208    "axvline": ("x",),
209    "axvspan": ("xmin", "xmax"),
210}
211
212
213def _convert_period_coords(axes: Axes, method_name: str, item: dict[str, Any]) -> dict[str, Any]:
214    """Return a copy of item with any Period x-coordinates replaced by ordinals.
215
216    If the axes was period-mapped by mgplot, the Period's freq must match the
217    axes' stashed freq — otherwise the ordinals live in different spaces.
218    On an axes with no stash we trust the programmer and just take .ordinal.
219    """
220    keys = _PERIOD_COORD_KEYS.get(method_name)
221    if not keys:
222        return item
223    stash = get_period_axes(axes)
224    stashed_freq = stash[0] if stash is not None else None
225    converted = dict(item)
226    for key in keys:
227        val = converted.get(key)
228        if isinstance(val, Period):
229            if stashed_freq is not None and val.freqstr != stashed_freq:
230                raise ValueError(
231                    f"{method_name} Period freq {val.freqstr!r} does not match axes freq {stashed_freq!r}",
232                )
233            if stashed_freq is not None:
234                # Widen the stash so later label refresh covers this coordinate,
235                # even if it falls outside the plotted data's ordinal range.
236                register_period_axes(axes, PeriodIndex([val]))
237            converted[key] = val.ordinal
238    return converted
239
240
241# Value-kwargs whose entries are x-axis coordinates — like axvline/axvspan, a
242# Period passed here must be converted to its ordinal on a period-mapped axes.
243_PERIOD_X_VALUE_KWARGS: Final[tuple[str, ...]] = ("xlim", "xticks")
244
245
246def _convert_period_value(axes: Axes, setting: str, value: Any) -> Any:
247    """Return value with any Period x-coordinates replaced by ordinals.
248
249    xlim is a 2-tuple and xticks a list; either may contain Periods when the
250    caller describes the axis in calendar terms. On a period-mapped axes the
251    Period freq must match the axes' stashed freq (see register_period_axes).
252    Non-x settings, None, and non-Period entries pass through unchanged.
253    """
254    if setting not in _PERIOD_X_VALUE_KWARGS or not isinstance(value, (tuple, list)):
255        return value
256    stash = get_period_axes(axes)
257    stashed_freq = stash[0] if stash is not None else None
258    converted: list[Any] = []
259    for val in value:
260        if isinstance(val, Period):
261            if stashed_freq is not None and val.freqstr != stashed_freq:
262                raise ValueError(
263                    f"{setting} Period freq {val.freqstr!r} does not match axes freq {stashed_freq!r}",
264                )
265            if stashed_freq is not None:
266                # Widen the stash so the later label refresh covers this coordinate.
267                register_period_axes(axes, PeriodIndex([val]))
268            converted.append(val.ordinal)
269        else:
270            converted.append(val)
271    return tuple(converted) if isinstance(value, tuple) else converted
272
273
274def _apply_splat(axes: Axes, method_name: str, value: _SplatValue) -> None:
275    """Apply a single splat kwarg, which may be a dict or sequence of dicts."""
276    if value is None or value is False:
277        return
278
279    if value is True:  # use the global default settings
280        value = get_setting(method_name)
281
282    # normalise to a list of dicts
283    if isinstance(value, dict):
284        value = [value]
285
286    if isinstance(value, Sequence):
287        method = getattr(axes, method_name)
288        for item in value:
289            if isinstance(item, dict):
290                method(**_convert_period_coords(axes, method_name, item))
291            else:
292                print(f"Warning: expected dict in {method_name} sequence, but got {type(item)}.")
293    else:
294        print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
295
296
297def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None:
298    """Set matplotlib elements dynamically using setting_name and splat."""
299    for method_name in settings:
300        if method_name not in kwargs:
301            continue
302
303        if method_name == "legend":
304            legend_value = kwargs.get(method_name)
305            if isinstance(legend_value, (bool, dict, type(None))):
306                make_legend(axes, legend=legend_value)
307            else:
308                print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.")
309            continue
310
311        value = kwargs.get(method_name)
312        if value is None or isinstance(value, (bool, dict, Sequence)):
313            _apply_splat(axes, method_name, value)
314        else:
315            print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
316
317
318def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
319    """Set figure size and apply chart annotations.
320
321    No-op when axes_only=True: the work here is all figure-level (resize,
322    corner text) and would stomp on other panels in a multi-axes figure.
323    """
324    if kwargs.get("axes_only"):
325        return
326    fig = axes.figure
327    fig_size = kwargs.get("figsize", get_setting("figsize"))
328    if not isinstance(fig, SubFigure):
329        fig.set_size_inches(*fig_size)
330
331    annotations = {
332        "rfooter": (0.99, 0.001, "right", "bottom"),
333        "lfooter": (0.01, 0.001, "left", "bottom"),
334        "rheader": (0.99, 0.999, "right", "top"),
335        "lheader": (0.01, 0.999, "left", "top"),
336    }
337
338    for annotation in HEADER_FOOTER_KWARGS:
339        if annotation in kwargs:
340            x_pos, y_pos, h_align, v_align = annotations[annotation]
341            fig.text(
342                x_pos,
343                y_pos,
344                str(kwargs.get(annotation, "")),
345                ha=h_align,
346                va=v_align,
347                fontsize=FOOTNOTE_FONTSIZE,
348                fontstyle=FOOTNOTE_FONTSTYLE,
349                color=FOOTNOTE_COLOR,
350            )
351
352
353def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
354    """Apply settings found in kwargs, after plotting the data."""
355    apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs)
356
357
358def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
359    """Apply settings found in kwargs."""
360
361    def check_kwargs(name: str) -> bool:
362        return name in kwargs and bool(kwargs.get(name))
363
364    apply_value_kwargs(axes, VALUE_KWARGS, **kwargs)
365    apply_annotations(axes, **kwargs)
366
367    if check_kwargs("zero_y"):
368        bottom, top = axes.get_ylim()
369        adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT
370        if bottom > -adj:
371            axes.set_ylim(bottom=-adj)
372        if top < adj:
373            axes.set_ylim(top=adj)
374
375    if check_kwargs("y0"):
376        low, high = axes.get_ylim()
377        if low < 0 < high:
378            axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
379
380    if check_kwargs("x0"):
381        low, high = axes.get_xlim()
382        if low < 0 < high:
383            axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
384
385    if check_kwargs("axisbelow"):
386        axes.set_axisbelow(True)
387
388
389def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
390    """Save the figure to file."""
391    saving = not kwargs.get("dont_save", False)  # save by default
392    if not saving:
393        return
394
395    try:
396        chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir")))
397
398        # Ensure directory exists
399        chart_dir.mkdir(parents=True, exist_ok=True)
400
401        suptitle = kwargs.get("suptitle", "")
402        title = kwargs.get("title", "")
403        pre_tag = kwargs.get("pre_tag", "")
404        tag = kwargs.get("tag", "")
405        name_override = kwargs.get("filename", "")
406        name_title = name_override or suptitle or title
407        file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME)
408        file_type = kwargs.get("file_type", get_setting("file_type")).lower()
409        dpi = kwargs.get("dpi", get_setting("dpi"))
410
411        # Construct filename components safely
412        filename_parts = []
413        if pre_tag:
414            filename_parts.append(sanitize_filename(pre_tag))
415        filename_parts.append(file_title)
416        if tag:
417            filename_parts.append(sanitize_filename(tag))
418
419        # Join filename parts and add extension
420        filename = "-".join(filter(None, filename_parts))
421        filepath = chart_dir / f"{filename}.{file_type}"
422
423        fig.savefig(filepath, dpi=dpi)
424
425    except (
426        OSError,
427        PermissionError,
428        FileNotFoundError,
429        ValueError,
430        RuntimeError,
431        TypeError,
432        UnicodeError,
433    ) as e:
434        print(f"Error: Could not save plot to file: {e}")
435
436
437# - public functions for finalise_plot()
438
439
440def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
441    """Finalise and save plots to the file system.
442
443    The filename for the saved plot is constructed from the global
444    chart_dir, the plot's title, any specified tag text, and the
445    file_type for the plot.
446
447    Args:
448        axes: Axes - matplotlib axes object - required
449        kwargs: FinaliseKwargs
450
451    """
452    # --- check the kwargs
453    report_kwargs(caller=ME, **kwargs)
454    validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs)
455
456    # --- sanity checks
457    if len(axes.get_children()) < 1:
458        print(f"Warning: {ME}() called with an empty axes, which was ignored.")
459        return
460
461    # --- remember axis-limits should we need to restore thems
462    xlim, ylim = axes.get_xlim(), axes.get_ylim()
463
464    # margins
465    axes.margins(DEFAULT_MARGIN)
466    axes.autoscale(tight=False)  # This is problematic ...
467
468    apply_kwargs(axes, **kwargs)
469
470    # tight layout and save the figure
471    fig = axes.figure
472    axes_only = kwargs.get("axes_only", False)
473    if not axes_only and (suptitle := kwargs.get("suptitle")):
474        fig.suptitle(suptitle)
475    if kwargs.get("preserve_lims"):
476        # restore the original limits of the axes
477        axes.set_xlim(xlim)
478        axes.set_ylim(ylim)
479    if not axes_only and not isinstance(fig, SubFigure):
480        fig.tight_layout(pad=TIGHT_LAYOUT_PAD)
481    apply_late_kwargs(axes, **kwargs)
482    # axvspan/axvline in late_kwargs may have widened xlim beyond what
483    # set_labels() last saw; regenerate ticks from the updated view.
484    refresh_period_labels(axes)
485    # de-collide end-of-line annotations now that the layout/limits are final
486    resolve_annotation_collisions(axes)
487    legend = axes.get_legend()
488    if legend and kwargs.get("remove_legend", False):
489        legend.remove()
490    if not axes_only and not isinstance(fig, SubFigure):
491        save_to_file(fig, **kwargs)
492
493    # show the plot in Jupyter Lab
494    if not axes_only and kwargs.get("show"):
495        plt.show()
496
497    # And close - the figure this axes belongs to, not pyplot's current figure
498    if not axes_only and not kwargs.get("dont_close", False):
499        root = fig
500        while isinstance(root, SubFigure):
501            root = root.figure
502        plt.close(root)