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, float] | None]
 42    ylim: NotRequired[tuple[float, float] | None]
 43    xticks: NotRequired[list[float] | 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 = 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
241def _apply_splat(axes: Axes, method_name: str, value: _SplatValue) -> None:
242    """Apply a single splat kwarg, which may be a dict or sequence of dicts."""
243    if value is None or value is False:
244        return
245
246    if value is True:  # use the global default settings
247        value = get_setting(method_name)
248
249    # normalise to a list of dicts
250    if isinstance(value, dict):
251        value = [value]
252
253    if isinstance(value, Sequence):
254        method = getattr(axes, method_name)
255        for item in value:
256            if isinstance(item, dict):
257                method(**_convert_period_coords(axes, method_name, item))
258            else:
259                print(f"Warning: expected dict in {method_name} sequence, but got {type(item)}.")
260    else:
261        print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
262
263
264def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None:
265    """Set matplotlib elements dynamically using setting_name and splat."""
266    for method_name in settings:
267        if method_name not in kwargs:
268            continue
269
270        if method_name == "legend":
271            legend_value = kwargs.get(method_name)
272            if isinstance(legend_value, (bool, dict, type(None))):
273                make_legend(axes, legend=legend_value)
274            else:
275                print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.")
276            continue
277
278        value = kwargs.get(method_name)
279        if value is None or isinstance(value, (bool, dict, Sequence)):
280            _apply_splat(axes, method_name, value)
281        else:
282            print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
283
284
285def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
286    """Set figure size and apply chart annotations.
287
288    No-op when axes_only=True: the work here is all figure-level (resize,
289    corner text) and would stomp on other panels in a multi-axes figure.
290    """
291    if kwargs.get("axes_only"):
292        return
293    fig = axes.figure
294    fig_size = kwargs.get("figsize", get_setting("figsize"))
295    if not isinstance(fig, SubFigure):
296        fig.set_size_inches(*fig_size)
297
298    annotations = {
299        "rfooter": (0.99, 0.001, "right", "bottom"),
300        "lfooter": (0.01, 0.001, "left", "bottom"),
301        "rheader": (0.99, 0.999, "right", "top"),
302        "lheader": (0.01, 0.999, "left", "top"),
303    }
304
305    for annotation in HEADER_FOOTER_KWARGS:
306        if annotation in kwargs:
307            x_pos, y_pos, h_align, v_align = annotations[annotation]
308            fig.text(
309                x_pos,
310                y_pos,
311                str(kwargs.get(annotation, "")),
312                ha=h_align,
313                va=v_align,
314                fontsize=FOOTNOTE_FONTSIZE,
315                fontstyle=FOOTNOTE_FONTSTYLE,
316                color=FOOTNOTE_COLOR,
317            )
318
319
320def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
321    """Apply settings found in kwargs, after plotting the data."""
322    apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs)
323
324
325def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
326    """Apply settings found in kwargs."""
327
328    def check_kwargs(name: str) -> bool:
329        return name in kwargs and bool(kwargs.get(name))
330
331    apply_value_kwargs(axes, VALUE_KWARGS, **kwargs)
332    apply_annotations(axes, **kwargs)
333
334    if check_kwargs("zero_y"):
335        bottom, top = axes.get_ylim()
336        adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT
337        if bottom > -adj:
338            axes.set_ylim(bottom=-adj)
339        if top < adj:
340            axes.set_ylim(top=adj)
341
342    if check_kwargs("y0"):
343        low, high = axes.get_ylim()
344        if low < 0 < high:
345            axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
346
347    if check_kwargs("x0"):
348        low, high = axes.get_xlim()
349        if low < 0 < high:
350            axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
351
352    if check_kwargs("axisbelow"):
353        axes.set_axisbelow(True)
354
355
356def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
357    """Save the figure to file."""
358    saving = not kwargs.get("dont_save", False)  # save by default
359    if not saving:
360        return
361
362    try:
363        chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir")))
364
365        # Ensure directory exists
366        chart_dir.mkdir(parents=True, exist_ok=True)
367
368        suptitle = kwargs.get("suptitle", "")
369        title = kwargs.get("title", "")
370        pre_tag = kwargs.get("pre_tag", "")
371        tag = kwargs.get("tag", "")
372        name_override = kwargs.get("filename", "")
373        name_title = name_override or suptitle or title
374        file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME)
375        file_type = kwargs.get("file_type", get_setting("file_type")).lower()
376        dpi = kwargs.get("dpi", get_setting("dpi"))
377
378        # Construct filename components safely
379        filename_parts = []
380        if pre_tag:
381            filename_parts.append(sanitize_filename(pre_tag))
382        filename_parts.append(file_title)
383        if tag:
384            filename_parts.append(sanitize_filename(tag))
385
386        # Join filename parts and add extension
387        filename = "-".join(filter(None, filename_parts))
388        filepath = chart_dir / f"{filename}.{file_type}"
389
390        fig.savefig(filepath, dpi=dpi)
391
392    except (
393        OSError,
394        PermissionError,
395        FileNotFoundError,
396        ValueError,
397        RuntimeError,
398        TypeError,
399        UnicodeError,
400    ) as e:
401        print(f"Error: Could not save plot to file: {e}")
402
403
404# - public functions for finalise_plot()
405
406
407def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
408    """Finalise and save plots to the file system.
409
410    The filename for the saved plot is constructed from the global
411    chart_dir, the plot's title, any specified tag text, and the
412    file_type for the plot.
413
414    Args:
415        axes: Axes - matplotlib axes object - required
416        kwargs: FinaliseKwargs
417
418    """
419    # --- check the kwargs
420    report_kwargs(caller=ME, **kwargs)
421    validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs)
422
423    # --- sanity checks
424    if len(axes.get_children()) < 1:
425        print(f"Warning: {ME}() called with an empty axes, which was ignored.")
426        return
427
428    # --- remember axis-limits should we need to restore thems
429    xlim, ylim = axes.get_xlim(), axes.get_ylim()
430
431    # margins
432    axes.margins(DEFAULT_MARGIN)
433    axes.autoscale(tight=False)  # This is problematic ...
434
435    apply_kwargs(axes, **kwargs)
436
437    # tight layout and save the figure
438    fig = axes.figure
439    axes_only = kwargs.get("axes_only", False)
440    if not axes_only and (suptitle := kwargs.get("suptitle")):
441        fig.suptitle(suptitle)
442    if kwargs.get("preserve_lims"):
443        # restore the original limits of the axes
444        axes.set_xlim(xlim)
445        axes.set_ylim(ylim)
446    if not axes_only and not isinstance(fig, SubFigure):
447        fig.tight_layout(pad=TIGHT_LAYOUT_PAD)
448    apply_late_kwargs(axes, **kwargs)
449    # axvspan/axvline in late_kwargs may have widened xlim beyond what
450    # set_labels() last saw; regenerate ticks from the updated view.
451    refresh_period_labels(axes)
452    # de-collide end-of-line annotations now that the layout/limits are final
453    resolve_annotation_collisions(axes)
454    legend = axes.get_legend()
455    if legend and kwargs.get("remove_legend", False):
456        legend.remove()
457    if not axes_only and not isinstance(fig, SubFigure):
458        save_to_file(fig, **kwargs)
459
460    # show the plot in Jupyter Lab
461    if not axes_only and kwargs.get("show"):
462        plt.show()
463
464    # And close - the figure this axes belongs to, not pyplot's current figure
465    if not axes_only and not kwargs.get("dont_close", False):
466        root = fig
467        while isinstance(root, SubFigure):
468            root = root.figure
469        plt.close(root)