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