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