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