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