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] | None]
 47    axvspan: NotRequired[dict[str, Any] | None]
 48    axhline: NotRequired[dict[str, Any] | None]
 49    axvline: NotRequired[dict[str, Any] | None]
 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            value = kwargs.get(method_name)
211            if value is None or value is False:
212                continue
213
214            if value is True:  # use the global default settings
215                value = get_setting(method_name)
216
217            # splat the kwargs to the method
218            if isinstance(value, dict):
219                method = getattr(axes, method_name)
220                method(**value)
221            else:
222                print(
223                    f"Warning expected dict argument for {method_name} but got {type(value)}.",
224                )
225
226
227def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
228    """Set figure size and apply chart annotations."""
229    fig = axes.figure
230    fig_size = kwargs.get("figsize", get_setting("figsize"))
231    if not isinstance(fig, SubFigure):
232        fig.set_size_inches(*fig_size)
233
234    annotations = {
235        "rfooter": (0.99, 0.001, "right", "bottom"),
236        "lfooter": (0.01, 0.001, "left", "bottom"),
237        "rheader": (0.99, 0.999, "right", "top"),
238        "lheader": (0.01, 0.999, "left", "top"),
239    }
240
241    for annotation in HEADER_FOOTER_KWARGS:
242        if annotation in kwargs:
243            x_pos, y_pos, h_align, v_align = annotations[annotation]
244            fig.text(
245                x_pos,
246                y_pos,
247                str(kwargs.get(annotation, "")),
248                ha=h_align,
249                va=v_align,
250                fontsize=FOOTNOTE_FONTSIZE,
251                fontstyle=FOOTNOTE_FONTSTYLE,
252                color=FOOTNOTE_COLOR,
253            )
254
255
256def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
257    """Apply settings found in kwargs, after plotting the data."""
258    apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs)
259
260
261def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
262    """Apply settings found in kwargs."""
263
264    def check_kwargs(name: str) -> bool:
265        return name in kwargs and bool(kwargs.get(name))
266
267    apply_value_kwargs(axes, VALUE_KWARGS, **kwargs)
268    apply_annotations(axes, **kwargs)
269
270    if check_kwargs("zero_y"):
271        bottom, top = axes.get_ylim()
272        adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT
273        if bottom > -adj:
274            axes.set_ylim(bottom=-adj)
275        if top < adj:
276            axes.set_ylim(top=adj)
277
278    if check_kwargs("y0"):
279        low, high = axes.get_ylim()
280        if low < 0 < high:
281            axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
282
283    if check_kwargs("x0"):
284        low, high = axes.get_xlim()
285        if low < 0 < high:
286            axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR)
287
288    if check_kwargs("axisbelow"):
289        axes.set_axisbelow(True)
290
291
292def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None:
293    """Save the figure to file."""
294    saving = not kwargs.get("dont_save", False)  # save by default
295    if not saving:
296        return
297
298    try:
299        chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir")))
300
301        # Ensure directory exists
302        chart_dir.mkdir(parents=True, exist_ok=True)
303
304        suptitle = kwargs.get("suptitle", "")
305        title = kwargs.get("title", "")
306        pre_tag = kwargs.get("pre_tag", "")
307        tag = kwargs.get("tag", "")
308        name_title = suptitle if suptitle else title
309        file_title = sanitize_filename(name_title if name_title else DEFAULT_FILE_TITLE_NAME)
310        file_type = kwargs.get("file_type", get_setting("file_type")).lower()
311        dpi = kwargs.get("dpi", get_setting("dpi"))
312
313        # Construct filename components safely
314        filename_parts = []
315        if pre_tag:
316            filename_parts.append(sanitize_filename(pre_tag))
317        filename_parts.append(file_title)
318        if tag:
319            filename_parts.append(sanitize_filename(tag))
320
321        # Join filename parts and add extension
322        filename = "-".join(filter(None, filename_parts))
323        filepath = chart_dir / f"{filename}.{file_type}"
324
325        fig.savefig(filepath, dpi=dpi)
326
327    except (
328        OSError,
329        PermissionError,
330        FileNotFoundError,
331        ValueError,
332        RuntimeError,
333        TypeError,
334        UnicodeError,
335    ) as e:
336        print(f"Error: Could not save plot to file: {e}")
337
338
339# - public functions for finalise_plot()
340
341
342def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
343    """Finalise and save plots to the file system.
344
345    The filename for the saved plot is constructed from the global
346    chart_dir, the plot's title, any specified tag text, and the
347    file_type for the plot.
348
349    Args:
350        axes: Axes - matplotlib axes object - required
351        kwargs: FinaliseKwargs
352
353    """
354    # --- check the kwargs
355    report_kwargs(caller=ME, **kwargs)
356    validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs)
357
358    # --- sanity checks
359    if len(axes.get_children()) < 1:
360        print(f"Warning: {ME}() called with an empty axes, which was ignored.")
361        return
362
363    # --- remember axis-limits should we need to restore thems
364    xlim, ylim = axes.get_xlim(), axes.get_ylim()
365
366    # margins
367    axes.margins(DEFAULT_MARGIN)
368    axes.autoscale(tight=False)  # This is problematic ...
369
370    apply_kwargs(axes, **kwargs)
371
372    # tight layout and save the figure
373    fig = axes.figure
374    if suptitle := kwargs.get("suptitle"):
375        fig.suptitle(suptitle)
376    if kwargs.get("preserve_lims"):
377        # restore the original limits of the axes
378        axes.set_xlim(xlim)
379        axes.set_ylim(ylim)
380    if not isinstance(fig, SubFigure):
381        fig.tight_layout(pad=TIGHT_LAYOUT_PAD)
382    apply_late_kwargs(axes, **kwargs)
383    legend = axes.get_legend()
384    if legend and kwargs.get("remove_legend", False):
385        legend.remove()
386    if not isinstance(fig, SubFigure):
387        save_to_file(fig, **kwargs)
388
389    # show the plot in Jupyter Lab
390    if kwargs.get("show"):
391        plt.show()
392
393    # And close
394    if not kwargs.get("dont_close", False):
395        plt.close()