mgplot.finalise_plot
Functions to finalise and save plots to the file system.
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 10import numpy as np 11from matplotlib.axes import Axes 12from matplotlib.figure import Figure, SubFigure 13from matplotlib.lines import Line2D 14from matplotlib.patches import Rectangle 15from matplotlib.transforms import blended_transform_factory 16from pandas import Period, PeriodIndex 17 18from mgplot.annotation_utils import resolve_annotation_collisions 19from mgplot.axis_utils import get_period_axes, refresh_period_labels, register_period_axes 20from mgplot.keyword_checking import BaseKwargs, report_kwargs, validate_kwargs 21from mgplot.settings import get_setting 22 23# --- constants 24ME: Final[str] = "finalise_plot" 25MAX_FILENAME_LENGTH: Final[int] = 150 26DEFAULT_MARGIN: Final[float] = 0.02 27TIGHT_LAYOUT_PAD: Final[float] = 1.1 28FOOTNOTE_FONTSIZE: Final[int] = 8 29FOOTNOTE_FONTSTYLE: Final[str] = "italic" 30FOOTNOTE_COLOR: Final[str] = "#999999" 31ZERO_LINE_WIDTH: Final[float] = 0.66 32ZERO_LINE_COLOR: Final[str] = "#555555" 33ZERO_AXIS_ADJUSTMENT: Final[float] = 0.02 34DEFAULT_FILE_TITLE_NAME: Final[str] = "plot" 35# --- annotated axvline text 36VLINE_TEXT_FONTSIZE: Final[str] = "xx-small" 37VLINE_TEXT_ROTATION: Final[int] = 90 38VLINE_TEXT_OFFSET: Final[float] = 3.0 # points, to the right of the line 39VLINE_TEXT_PAD: Final[float] = 0.01 # axes fraction, in from the top/bottom 40VLINE_AUTO_BAND: Final[float] = 0.02 # fraction of the x-span sampled either side 41 42 43class FinaliseKwargs(BaseKwargs): 44 """Keyword arguments for the finalise_plot function.""" 45 46 # --- value options 47 suptitle: NotRequired[str | None] 48 title: NotRequired[str | None] 49 xlabel: NotRequired[str | None] 50 ylabel: NotRequired[str | None] 51 xlim: NotRequired[tuple[float | int | Period, float | int | Period] | None] 52 ylim: NotRequired[tuple[float, float] | None] 53 xticks: NotRequired[list[float | int | Period] | None] 54 yticks: NotRequired[list[float] | None] 55 xscale: NotRequired[str | None] 56 yscale: NotRequired[str | None] 57 # --- splat options 58 legend: NotRequired[bool | dict[str, Any] | None] 59 axhspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None] 60 axvspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None] 61 axhline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None] 62 axvline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None] 63 # --- options for annotations 64 lfooter: NotRequired[str] 65 rfooter: NotRequired[str] 66 lheader: NotRequired[str] 67 rheader: NotRequired[str] 68 # --- file/save options 69 pre_tag: NotRequired[str] 70 tag: NotRequired[str] 71 filename: NotRequired[str] 72 chart_dir: NotRequired[str] 73 file_type: NotRequired[str] 74 dpi: NotRequired[int] 75 figsize: NotRequired[tuple[float, float]] 76 show: NotRequired[bool] 77 # --- other options 78 preserve_lims: NotRequired[bool] 79 remove_legend: NotRequired[bool] 80 zero_y: NotRequired[bool] 81 y0: NotRequired[bool] 82 x0: NotRequired[bool] 83 axisbelow: NotRequired[bool] 84 dont_save: NotRequired[bool] 85 dont_close: NotRequired[bool] 86 axes_only: NotRequired[bool] 87 88 89VALUE_KWARGS = ( 90 "title", 91 "xlabel", 92 "ylabel", 93 "xlim", 94 "ylim", 95 "xticks", 96 "yticks", 97 "xscale", 98 "yscale", 99) 100SPLAT_KWARGS = ( 101 "axhspan", 102 "axvspan", 103 "axhline", 104 "axvline", 105 "legend", # needs to be last in this tuple 106) 107HEADER_FOOTER_KWARGS = ( 108 "lfooter", 109 "rfooter", 110 "lheader", 111 "rheader", 112) 113 114 115def sanitize_filename(filename: str, max_length: int = MAX_FILENAME_LENGTH) -> str: 116 """Convert a string to a safe filename. 117 118 Args: 119 filename: The string to convert to a filename 120 max_length: Maximum length for the filename 121 122 Returns: 123 A safe filename string 124 125 """ 126 if not filename: 127 return "untitled" 128 129 # Normalize unicode characters (e.g., é -> e) 130 filename = unicodedata.normalize("NFKD", filename) 131 132 # Remove non-ASCII characters 133 filename = filename.encode("ascii", "ignore").decode("ascii") 134 135 # Convert to lowercase 136 filename = filename.lower() 137 138 # Replace spaces and other separators with hyphens 139 filename = re.sub(r"[\s\-_]+", "-", filename) 140 141 # Remove unsafe characters, keeping only alphanumeric and hyphens 142 filename = re.sub(r"[^a-z0-9\-]", "", filename) 143 144 # Remove leading/trailing hyphens and collapse multiple hyphens 145 filename = re.sub(r"^-+|-+$", "", filename) 146 filename = re.sub(r"-+", "-", filename) 147 148 # Truncate to max length 149 if len(filename) > max_length: 150 filename = filename[:max_length].rstrip("-") 151 152 # Ensure we have a valid filename 153 return filename or "untitled" 154 155 156def make_legend(axes: Axes, *, legend: None | bool | dict[str, Any]) -> None: 157 """Create a legend for the plot.""" 158 if legend is None or legend is False: 159 return 160 161 if legend is True: # use the global default settings 162 legend = get_setting("legend") 163 164 if isinstance(legend, dict): 165 axes.legend(**legend) 166 return 167 168 print(f"Warning: expected dict argument for legend, but got {type(legend)}.") 169 170 171def apply_value_kwargs(axes: Axes, value_kwargs_: Sequence[str], **kwargs: Unpack[FinaliseKwargs]) -> None: 172 """Set matplotlib elements by name using Axes.set(). 173 174 Tricky: some plotting functions may set the xlabel or ylabel. 175 So ... we will set these if a setting is explicitly provided. If no 176 setting is provided, we will set to None if they are not already set. 177 If they have already been set, we will not change them. 178 179 """ 180 # --- preliminary 181 function: dict[str, Callable[[], str]] = { 182 "xlabel": axes.get_xlabel, 183 "ylabel": axes.get_ylabel, 184 "title": axes.get_title, 185 } 186 187 def fail() -> str: 188 return "" 189 190 # --- loop over potential value settings 191 for setting in value_kwargs_: 192 value = _convert_period_value(axes, setting, kwargs.get(setting)) 193 if setting in kwargs: 194 # deliberately set, so we will action 195 axes.set(**{setting: value}) 196 continue 197 required_to_set = ("title", "xlabel", "ylabel") 198 if setting not in required_to_set: 199 # not set - and not required - so we can skip 200 continue 201 202 # we will set these 'required_to_set' ones 203 # provided they are not already set 204 already_set = function.get(setting, fail)() 205 if already_set and value is None: 206 continue 207 208 # if we get here, we will set the value (implicitly to None) 209 axes.set(**{setting: value}) 210 211 212_SplatValue = bool | dict[str, Any] | Sequence[dict[str, Any]] | None 213 214# Keys in each splat-method's kwargs that are x-axis coordinates — when the 215# plot uses a PeriodIndex the axis is mapped to Period ordinals, so a Period 216# passed here must be converted to its ordinal for matplotlib. 217_PERIOD_COORD_KEYS: Final[dict[str, tuple[str, ...]]] = { 218 "axvline": ("x",), 219 "axvspan": ("xmin", "xmax"), 220} 221 222 223def _convert_period_coords(axes: Axes, method_name: str, item: dict[str, Any]) -> dict[str, Any]: 224 """Return a copy of item with any Period x-coordinates replaced by ordinals. 225 226 If the axes was period-mapped by mgplot, the Period's freq must match the 227 axes' stashed freq — otherwise the ordinals live in different spaces. 228 On an axes with no stash we trust the programmer and just take .ordinal. 229 """ 230 keys = _PERIOD_COORD_KEYS.get(method_name) 231 if not keys: 232 return item 233 stash = get_period_axes(axes) 234 stashed_freq = stash[0] if stash is not None else None 235 converted = dict(item) 236 for key in keys: 237 val = converted.get(key) 238 if isinstance(val, Period): 239 if stashed_freq is not None and val.freqstr != stashed_freq: 240 raise ValueError( 241 f"{method_name} Period freq {val.freqstr!r} does not match axes freq {stashed_freq!r}", 242 ) 243 if stashed_freq is not None: 244 # Widen the stash so later label refresh covers this coordinate, 245 # even if it falls outside the plotted data's ordinal range. 246 register_period_axes(axes, PeriodIndex([val])) 247 converted[key] = val.ordinal 248 return converted 249 250 251# Value-kwargs whose entries are x-axis coordinates — like axvline/axvspan, a 252# Period passed here must be converted to its ordinal on a period-mapped axes. 253_PERIOD_X_VALUE_KWARGS: Final[tuple[str, ...]] = ("xlim", "xticks") 254 255 256def _convert_period_value(axes: Axes, setting: str, value: object) -> object: 257 """Return value with any Period x-coordinates replaced by ordinals. 258 259 xlim is a 2-tuple and xticks a list; either may contain Periods when the 260 caller describes the axis in calendar terms. On a period-mapped axes the 261 Period freq must match the axes' stashed freq (see register_period_axes). 262 Non-x settings, None, and non-Period entries pass through unchanged. 263 """ 264 if setting not in _PERIOD_X_VALUE_KWARGS or not isinstance(value, (tuple, list)): 265 return value 266 stash = get_period_axes(axes) 267 stashed_freq = stash[0] if stash is not None else None 268 converted: list[object] = [] 269 for val in value: 270 if isinstance(val, Period): 271 if stashed_freq is not None and val.freqstr != stashed_freq: 272 raise ValueError( 273 f"{setting} Period freq {val.freqstr!r} does not match axes freq {stashed_freq!r}", 274 ) 275 if stashed_freq is not None: 276 # Widen the stash so the later label refresh covers this coordinate. 277 register_period_axes(axes, PeriodIndex([val])) 278 converted.append(val.ordinal) 279 else: 280 converted.append(val) 281 return tuple(converted) if isinstance(value, tuple) else converted 282 283 284# --- annotated vertical lines 285# "text", "loc" and "text_kwargs" in an axvline dict describe its label rather 286# than the line, and are popped before the dict is splatted into ax.axvline(). 287_VLINE_LOCATIONS: Final[tuple[str, ...]] = ("auto", "top", "bottom") 288 289 290def _pop_vline_text(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]] | None: 291 """Remove the label keys from an axvline dict and return them, or None if unlabelled. 292 293 Raises on a bad loc, a non-dict text_kwargs, or label options given 294 without any text to place -- each of which is a typo, not a choice. 295 """ 296 text = item.pop("text", None) 297 loc = item.pop("loc", "auto") 298 text_kwargs = item.pop("text_kwargs", {}) 299 300 if text is None or not str(text).strip(): 301 if loc != "auto" or text_kwargs: 302 raise ValueError("axvline 'loc'/'text_kwargs' given without any 'text' to place") 303 return None 304 if loc not in _VLINE_LOCATIONS: 305 raise ValueError(f"axvline 'loc' must be one of {_VLINE_LOCATIONS}, got {loc!r}") 306 if not isinstance(text_kwargs, dict): 307 raise TypeError(f"axvline 'text_kwargs' must be a dict, got {type(text_kwargs)}") 308 return str(text), loc, text_kwargs 309 310 311def _data_y_extent(axes: Axes, x: float) -> tuple[float, float] | None: 312 """Return the (min, max) y of plotted data in a narrow x-band around x. 313 314 A rotated label occupies only a sliver of the x-axis, so a narrow band is 315 what the eye actually judges. Lines are matched by transform identity: 316 axhline/axvline use blended transforms, so this skips them and measures 317 only real data. Rectangles cover bar plots; axvspan/axhspan patches are 318 likewise blended and skipped. Returns None when nothing is measurable. 319 """ 320 left, right = axes.get_xlim() 321 half = abs(right - left) * VLINE_AUTO_BAND 322 low, high = x - half, x + half 323 found: list[float] = [] 324 325 for line in axes.get_lines(): 326 if not isinstance(line, Line2D) or line.get_transform() is not axes.transData: 327 continue 328 xdata = np.asarray(line.get_xdata(), dtype=float) 329 ydata = np.asarray(line.get_ydata(), dtype=float) 330 if xdata.size != ydata.size or xdata.size == 0: 331 continue 332 wanted = (xdata >= low) & (xdata <= high) & ~np.isnan(ydata) 333 if wanted.any(): 334 found.extend((float(ydata[wanted].min()), float(ydata[wanted].max()))) 335 336 for patch in axes.patches: 337 if not isinstance(patch, Rectangle) or patch.get_data_transform() is not axes.transData: 338 continue 339 x0, width = patch.get_x(), patch.get_width() 340 y0, height = patch.get_y(), patch.get_height() 341 if min(x0, x0 + width) > high or max(x0, x0 + width) < low: 342 continue 343 found.extend((min(y0, y0 + height), max(y0, y0 + height))) 344 345 if not found: 346 return None 347 return min(found), max(found) 348 349 350def _auto_vline_loc(axes: Axes, x: float) -> str: 351 """Pick the end of the axes with more room between the data and the limit.""" 352 extent = _data_y_extent(axes, x) 353 if extent is None: 354 return "top" # nothing measurable (empty band, or an unrecognised artist) 355 data_low, data_high = extent 356 bottom_lim, top_lim = axes.get_ylim() 357 if top_lim >= bottom_lim: 358 top_gap, bottom_gap = top_lim - data_high, data_low - bottom_lim 359 else: # inverted y-axis: values decrease going up the display 360 top_gap, bottom_gap = data_low - top_lim, bottom_lim - data_high 361 return "top" if top_gap >= bottom_gap else "bottom" 362 363 364def _annotate_vline(axes: Axes, item: dict[str, Any], line: Line2D, spec: tuple[str, str, dict]) -> None: 365 """Place a rotated text label just to the right of a vertical line. 366 367 The label is anchored with x in data coordinates and y in axes 368 coordinates, so it stays pinned to the top/bottom of the plot regardless 369 of any later change to the y-limits. 370 """ 371 text, loc, text_kwargs = spec 372 x = item.get("x", 0) # matches the matplotlib default for axvline 373 if loc == "auto": 374 loc = _auto_vline_loc(axes, float(x)) 375 y, valign = (1.0 - VLINE_TEXT_PAD, "top") if loc == "top" else (VLINE_TEXT_PAD, "bottom") 376 377 options: dict[str, Any] = { 378 "rotation": VLINE_TEXT_ROTATION, 379 "fontsize": VLINE_TEXT_FONTSIZE, 380 "color": line.get_color(), 381 "ha": "left", 382 "va": valign, 383 } 384 options.update(text_kwargs) 385 axes.annotate( 386 text, 387 xy=(x, y), 388 xycoords=blended_transform_factory(axes.transData, axes.transAxes), 389 xytext=(VLINE_TEXT_OFFSET, 0), 390 textcoords="offset points", 391 **options, 392 ) 393 394 395def _apply_splat(axes: Axes, method_name: str, value: _SplatValue) -> None: 396 """Apply a single splat kwarg, which may be a dict or sequence of dicts.""" 397 if value is None or value is False: 398 return 399 400 if value is True: # use the global default settings 401 value = get_setting(method_name) 402 403 # normalise to a list of dicts 404 if isinstance(value, dict): 405 value = [value] 406 407 if isinstance(value, Sequence): 408 method = getattr(axes, method_name) 409 for item in value: 410 if not isinstance(item, dict): 411 print(f"Warning: expected dict in {method_name} sequence, but got {type(item)}.") 412 continue 413 converted = _convert_period_coords(axes, method_name, item) 414 # _convert_period_coords always copies, so popping is safe here 415 spec = _pop_vline_text(converted) if method_name == "axvline" else None 416 artist = method(**converted) 417 if spec is not None and isinstance(artist, Line2D): 418 _annotate_vline(axes, converted, artist, spec) 419 else: 420 print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.") 421 422 423def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None: 424 """Set matplotlib elements dynamically using setting_name and splat.""" 425 for method_name in settings: 426 if method_name not in kwargs: 427 continue 428 429 if method_name == "legend": 430 legend_value = kwargs.get(method_name) 431 if isinstance(legend_value, (bool, dict, type(None))): 432 make_legend(axes, legend=legend_value) 433 else: 434 print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.") 435 continue 436 437 value = kwargs.get(method_name) 438 if value is None or isinstance(value, (bool, dict, Sequence)): 439 _apply_splat(axes, method_name, value) 440 else: 441 print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.") 442 443 444def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 445 """Set figure size and apply chart annotations. 446 447 No-op when axes_only=True: the work here is all figure-level (resize, 448 corner text) and would stomp on other panels in a multi-axes figure. 449 """ 450 if kwargs.get("axes_only"): 451 return 452 fig = axes.figure 453 fig_size = kwargs.get("figsize", get_setting("figsize")) 454 if not isinstance(fig, SubFigure): 455 fig.set_size_inches(*fig_size) 456 457 annotations = { 458 "rfooter": (0.99, 0.001, "right", "bottom"), 459 "lfooter": (0.01, 0.001, "left", "bottom"), 460 "rheader": (0.99, 0.999, "right", "top"), 461 "lheader": (0.01, 0.999, "left", "top"), 462 } 463 464 for annotation in HEADER_FOOTER_KWARGS: 465 if annotation in kwargs: 466 x_pos, y_pos, h_align, v_align = annotations[annotation] 467 fig.text( 468 x_pos, 469 y_pos, 470 str(kwargs.get(annotation, "")), 471 ha=h_align, 472 va=v_align, 473 fontsize=FOOTNOTE_FONTSIZE, 474 fontstyle=FOOTNOTE_FONTSTYLE, 475 color=FOOTNOTE_COLOR, 476 ) 477 478 479def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 480 """Apply settings found in kwargs, after plotting the data.""" 481 apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs) 482 483 484def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 485 """Apply settings found in kwargs.""" 486 487 def check_kwargs(name: str) -> bool: 488 return name in kwargs and bool(kwargs.get(name)) 489 490 apply_value_kwargs(axes, VALUE_KWARGS, **kwargs) 491 apply_annotations(axes, **kwargs) 492 493 if check_kwargs("zero_y"): 494 bottom, top = axes.get_ylim() 495 adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT 496 if bottom > -adj: 497 axes.set_ylim(bottom=-adj) 498 if top < adj: 499 axes.set_ylim(top=adj) 500 501 if check_kwargs("y0"): 502 low, high = axes.get_ylim() 503 if low < 0 < high: 504 axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR) 505 506 if check_kwargs("x0"): 507 low, high = axes.get_xlim() 508 if low < 0 < high: 509 axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR) 510 511 if check_kwargs("axisbelow"): 512 axes.set_axisbelow(True) 513 514 515def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None: 516 """Save the figure to file.""" 517 saving = not kwargs.get("dont_save", False) # save by default 518 if not saving: 519 return 520 521 try: 522 chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir"))) 523 524 # Ensure directory exists 525 chart_dir.mkdir(parents=True, exist_ok=True) 526 527 suptitle = kwargs.get("suptitle", "") 528 title = kwargs.get("title", "") 529 pre_tag = kwargs.get("pre_tag", "") 530 tag = kwargs.get("tag", "") 531 name_override = kwargs.get("filename", "") 532 name_title = name_override or suptitle or title 533 file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME) 534 file_type = kwargs.get("file_type", get_setting("file_type")).lower() 535 dpi = kwargs.get("dpi", get_setting("dpi")) 536 537 # Construct filename components safely 538 filename_parts = [] 539 if pre_tag: 540 filename_parts.append(sanitize_filename(pre_tag)) 541 filename_parts.append(file_title) 542 if tag: 543 filename_parts.append(sanitize_filename(tag)) 544 545 # Join filename parts and add extension 546 filename = "-".join(filter(None, filename_parts)) 547 filepath = chart_dir / f"{filename}.{file_type}" 548 549 fig.savefig(filepath, dpi=dpi) 550 551 except ( 552 OSError, 553 PermissionError, 554 FileNotFoundError, 555 ValueError, 556 RuntimeError, 557 TypeError, 558 UnicodeError, 559 ) as e: 560 print(f"Error: Could not save plot to file: {e}") 561 562 563# - public functions for finalise_plot() 564 565 566def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 567 """Finalise and save plots to the file system. 568 569 The filename for the saved plot is constructed from the global 570 chart_dir, the plot's title, any specified tag text, and the 571 file_type for the plot. 572 573 Args: 574 axes: Axes - matplotlib axes object - required 575 kwargs: FinaliseKwargs 576 577 """ 578 # --- check the kwargs 579 report_kwargs(caller=ME, **kwargs) 580 validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs) 581 582 # --- sanity checks 583 if len(axes.get_children()) < 1: 584 print(f"Warning: {ME}() called with an empty axes, which was ignored.") 585 return 586 587 # --- remember axis-limits should we need to restore thems 588 xlim, ylim = axes.get_xlim(), axes.get_ylim() 589 590 # margins 591 axes.margins(DEFAULT_MARGIN) 592 axes.autoscale(tight=False) # This is problematic ... 593 594 apply_kwargs(axes, **kwargs) 595 596 # tight layout and save the figure 597 fig = axes.figure 598 axes_only = kwargs.get("axes_only", False) 599 if not axes_only and (suptitle := kwargs.get("suptitle")): 600 fig.suptitle(suptitle) 601 if kwargs.get("preserve_lims"): 602 # restore the original limits of the axes 603 axes.set_xlim(xlim) 604 axes.set_ylim(ylim) 605 if not axes_only and not isinstance(fig, SubFigure): 606 fig.tight_layout(pad=TIGHT_LAYOUT_PAD) 607 apply_late_kwargs(axes, **kwargs) 608 # axvspan/axvline in late_kwargs may have widened xlim beyond what 609 # set_labels() last saw; regenerate ticks from the updated view. 610 refresh_period_labels(axes) 611 # de-collide end-of-line annotations now that the layout/limits are final 612 resolve_annotation_collisions(axes) 613 legend = axes.get_legend() 614 if legend and kwargs.get("remove_legend", False): 615 legend.remove() 616 if not axes_only and not isinstance(fig, SubFigure): 617 save_to_file(fig, **kwargs) 618 619 # show the plot in Jupyter Lab 620 if not axes_only and kwargs.get("show"): 621 plt.show() 622 623 # And close - the figure this axes belongs to, not pyplot's current figure 624 if not axes_only and not kwargs.get("dont_close", False): 625 root = fig 626 while isinstance(root, SubFigure): 627 root = root.figure 628 plt.close(root)
44class FinaliseKwargs(BaseKwargs): 45 """Keyword arguments for the finalise_plot function.""" 46 47 # --- value options 48 suptitle: NotRequired[str | None] 49 title: NotRequired[str | None] 50 xlabel: NotRequired[str | None] 51 ylabel: NotRequired[str | None] 52 xlim: NotRequired[tuple[float | int | Period, float | int | Period] | None] 53 ylim: NotRequired[tuple[float, float] | None] 54 xticks: NotRequired[list[float | int | Period] | None] 55 yticks: NotRequired[list[float] | None] 56 xscale: NotRequired[str | None] 57 yscale: NotRequired[str | None] 58 # --- splat options 59 legend: NotRequired[bool | dict[str, Any] | None] 60 axhspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None] 61 axvspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None] 62 axhline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None] 63 axvline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None] 64 # --- options for annotations 65 lfooter: NotRequired[str] 66 rfooter: NotRequired[str] 67 lheader: NotRequired[str] 68 rheader: NotRequired[str] 69 # --- file/save options 70 pre_tag: NotRequired[str] 71 tag: NotRequired[str] 72 filename: NotRequired[str] 73 chart_dir: NotRequired[str] 74 file_type: NotRequired[str] 75 dpi: NotRequired[int] 76 figsize: NotRequired[tuple[float, float]] 77 show: NotRequired[bool] 78 # --- other options 79 preserve_lims: NotRequired[bool] 80 remove_legend: NotRequired[bool] 81 zero_y: NotRequired[bool] 82 y0: NotRequired[bool] 83 x0: NotRequired[bool] 84 axisbelow: NotRequired[bool] 85 dont_save: NotRequired[bool] 86 dont_close: NotRequired[bool] 87 axes_only: NotRequired[bool]
Keyword arguments for the finalise_plot function.
116def sanitize_filename(filename: str, max_length: int = MAX_FILENAME_LENGTH) -> str: 117 """Convert a string to a safe filename. 118 119 Args: 120 filename: The string to convert to a filename 121 max_length: Maximum length for the filename 122 123 Returns: 124 A safe filename string 125 126 """ 127 if not filename: 128 return "untitled" 129 130 # Normalize unicode characters (e.g., é -> e) 131 filename = unicodedata.normalize("NFKD", filename) 132 133 # Remove non-ASCII characters 134 filename = filename.encode("ascii", "ignore").decode("ascii") 135 136 # Convert to lowercase 137 filename = filename.lower() 138 139 # Replace spaces and other separators with hyphens 140 filename = re.sub(r"[\s\-_]+", "-", filename) 141 142 # Remove unsafe characters, keeping only alphanumeric and hyphens 143 filename = re.sub(r"[^a-z0-9\-]", "", filename) 144 145 # Remove leading/trailing hyphens and collapse multiple hyphens 146 filename = re.sub(r"^-+|-+$", "", filename) 147 filename = re.sub(r"-+", "-", filename) 148 149 # Truncate to max length 150 if len(filename) > max_length: 151 filename = filename[:max_length].rstrip("-") 152 153 # Ensure we have a valid filename 154 return filename or "untitled"
Convert a string to a safe filename.
Args: filename: The string to convert to a filename max_length: Maximum length for the filename
Returns: A safe filename string
157def make_legend(axes: Axes, *, legend: None | bool | dict[str, Any]) -> None: 158 """Create a legend for the plot.""" 159 if legend is None or legend is False: 160 return 161 162 if legend is True: # use the global default settings 163 legend = get_setting("legend") 164 165 if isinstance(legend, dict): 166 axes.legend(**legend) 167 return 168 169 print(f"Warning: expected dict argument for legend, but got {type(legend)}.")
Create a legend for the plot.
172def apply_value_kwargs(axes: Axes, value_kwargs_: Sequence[str], **kwargs: Unpack[FinaliseKwargs]) -> None: 173 """Set matplotlib elements by name using Axes.set(). 174 175 Tricky: some plotting functions may set the xlabel or ylabel. 176 So ... we will set these if a setting is explicitly provided. If no 177 setting is provided, we will set to None if they are not already set. 178 If they have already been set, we will not change them. 179 180 """ 181 # --- preliminary 182 function: dict[str, Callable[[], str]] = { 183 "xlabel": axes.get_xlabel, 184 "ylabel": axes.get_ylabel, 185 "title": axes.get_title, 186 } 187 188 def fail() -> str: 189 return "" 190 191 # --- loop over potential value settings 192 for setting in value_kwargs_: 193 value = _convert_period_value(axes, setting, kwargs.get(setting)) 194 if setting in kwargs: 195 # deliberately set, so we will action 196 axes.set(**{setting: value}) 197 continue 198 required_to_set = ("title", "xlabel", "ylabel") 199 if setting not in required_to_set: 200 # not set - and not required - so we can skip 201 continue 202 203 # we will set these 'required_to_set' ones 204 # provided they are not already set 205 already_set = function.get(setting, fail)() 206 if already_set and value is None: 207 continue 208 209 # if we get here, we will set the value (implicitly to None) 210 axes.set(**{setting: value})
Set matplotlib elements by name using Axes.set().
Tricky: some plotting functions may set the xlabel or ylabel. So ... we will set these if a setting is explicitly provided. If no setting is provided, we will set to None if they are not already set. If they have already been set, we will not change them.
424def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None: 425 """Set matplotlib elements dynamically using setting_name and splat.""" 426 for method_name in settings: 427 if method_name not in kwargs: 428 continue 429 430 if method_name == "legend": 431 legend_value = kwargs.get(method_name) 432 if isinstance(legend_value, (bool, dict, type(None))): 433 make_legend(axes, legend=legend_value) 434 else: 435 print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.") 436 continue 437 438 value = kwargs.get(method_name) 439 if value is None or isinstance(value, (bool, dict, Sequence)): 440 _apply_splat(axes, method_name, value) 441 else: 442 print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.")
Set matplotlib elements dynamically using setting_name and splat.
445def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 446 """Set figure size and apply chart annotations. 447 448 No-op when axes_only=True: the work here is all figure-level (resize, 449 corner text) and would stomp on other panels in a multi-axes figure. 450 """ 451 if kwargs.get("axes_only"): 452 return 453 fig = axes.figure 454 fig_size = kwargs.get("figsize", get_setting("figsize")) 455 if not isinstance(fig, SubFigure): 456 fig.set_size_inches(*fig_size) 457 458 annotations = { 459 "rfooter": (0.99, 0.001, "right", "bottom"), 460 "lfooter": (0.01, 0.001, "left", "bottom"), 461 "rheader": (0.99, 0.999, "right", "top"), 462 "lheader": (0.01, 0.999, "left", "top"), 463 } 464 465 for annotation in HEADER_FOOTER_KWARGS: 466 if annotation in kwargs: 467 x_pos, y_pos, h_align, v_align = annotations[annotation] 468 fig.text( 469 x_pos, 470 y_pos, 471 str(kwargs.get(annotation, "")), 472 ha=h_align, 473 va=v_align, 474 fontsize=FOOTNOTE_FONTSIZE, 475 fontstyle=FOOTNOTE_FONTSTYLE, 476 color=FOOTNOTE_COLOR, 477 )
Set figure size and apply chart annotations.
No-op when axes_only=True: the work here is all figure-level (resize, corner text) and would stomp on other panels in a multi-axes figure.
480def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 481 """Apply settings found in kwargs, after plotting the data.""" 482 apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs)
Apply settings found in kwargs, after plotting the data.
485def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 486 """Apply settings found in kwargs.""" 487 488 def check_kwargs(name: str) -> bool: 489 return name in kwargs and bool(kwargs.get(name)) 490 491 apply_value_kwargs(axes, VALUE_KWARGS, **kwargs) 492 apply_annotations(axes, **kwargs) 493 494 if check_kwargs("zero_y"): 495 bottom, top = axes.get_ylim() 496 adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT 497 if bottom > -adj: 498 axes.set_ylim(bottom=-adj) 499 if top < adj: 500 axes.set_ylim(top=adj) 501 502 if check_kwargs("y0"): 503 low, high = axes.get_ylim() 504 if low < 0 < high: 505 axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR) 506 507 if check_kwargs("x0"): 508 low, high = axes.get_xlim() 509 if low < 0 < high: 510 axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR) 511 512 if check_kwargs("axisbelow"): 513 axes.set_axisbelow(True)
Apply settings found in kwargs.
516def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None: 517 """Save the figure to file.""" 518 saving = not kwargs.get("dont_save", False) # save by default 519 if not saving: 520 return 521 522 try: 523 chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir"))) 524 525 # Ensure directory exists 526 chart_dir.mkdir(parents=True, exist_ok=True) 527 528 suptitle = kwargs.get("suptitle", "") 529 title = kwargs.get("title", "") 530 pre_tag = kwargs.get("pre_tag", "") 531 tag = kwargs.get("tag", "") 532 name_override = kwargs.get("filename", "") 533 name_title = name_override or suptitle or title 534 file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME) 535 file_type = kwargs.get("file_type", get_setting("file_type")).lower() 536 dpi = kwargs.get("dpi", get_setting("dpi")) 537 538 # Construct filename components safely 539 filename_parts = [] 540 if pre_tag: 541 filename_parts.append(sanitize_filename(pre_tag)) 542 filename_parts.append(file_title) 543 if tag: 544 filename_parts.append(sanitize_filename(tag)) 545 546 # Join filename parts and add extension 547 filename = "-".join(filter(None, filename_parts)) 548 filepath = chart_dir / f"{filename}.{file_type}" 549 550 fig.savefig(filepath, dpi=dpi) 551 552 except ( 553 OSError, 554 PermissionError, 555 FileNotFoundError, 556 ValueError, 557 RuntimeError, 558 TypeError, 559 UnicodeError, 560 ) as e: 561 print(f"Error: Could not save plot to file: {e}")
Save the figure to file.
567def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 568 """Finalise and save plots to the file system. 569 570 The filename for the saved plot is constructed from the global 571 chart_dir, the plot's title, any specified tag text, and the 572 file_type for the plot. 573 574 Args: 575 axes: Axes - matplotlib axes object - required 576 kwargs: FinaliseKwargs 577 578 """ 579 # --- check the kwargs 580 report_kwargs(caller=ME, **kwargs) 581 validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs) 582 583 # --- sanity checks 584 if len(axes.get_children()) < 1: 585 print(f"Warning: {ME}() called with an empty axes, which was ignored.") 586 return 587 588 # --- remember axis-limits should we need to restore thems 589 xlim, ylim = axes.get_xlim(), axes.get_ylim() 590 591 # margins 592 axes.margins(DEFAULT_MARGIN) 593 axes.autoscale(tight=False) # This is problematic ... 594 595 apply_kwargs(axes, **kwargs) 596 597 # tight layout and save the figure 598 fig = axes.figure 599 axes_only = kwargs.get("axes_only", False) 600 if not axes_only and (suptitle := kwargs.get("suptitle")): 601 fig.suptitle(suptitle) 602 if kwargs.get("preserve_lims"): 603 # restore the original limits of the axes 604 axes.set_xlim(xlim) 605 axes.set_ylim(ylim) 606 if not axes_only and not isinstance(fig, SubFigure): 607 fig.tight_layout(pad=TIGHT_LAYOUT_PAD) 608 apply_late_kwargs(axes, **kwargs) 609 # axvspan/axvline in late_kwargs may have widened xlim beyond what 610 # set_labels() last saw; regenerate ticks from the updated view. 611 refresh_period_labels(axes) 612 # de-collide end-of-line annotations now that the layout/limits are final 613 resolve_annotation_collisions(axes) 614 legend = axes.get_legend() 615 if legend and kwargs.get("remove_legend", False): 616 legend.remove() 617 if not axes_only and not isinstance(fig, SubFigure): 618 save_to_file(fig, **kwargs) 619 620 # show the plot in Jupyter Lab 621 if not axes_only and kwargs.get("show"): 622 plt.show() 623 624 # And close - the figure this axes belongs to, not pyplot's current figure 625 if not axes_only and not kwargs.get("dont_close", False): 626 root = fig 627 while isinstance(root, SubFigure): 628 root = root.figure 629 plt.close(root)
Finalise and save plots to the file system.
The filename for the saved plot is constructed from the global chart_dir, the plot's title, any specified tag text, and the file_type for the plot.
Args: axes: Axes - matplotlib axes object - required kwargs: FinaliseKwargs