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