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] = 2.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# loc is a space-separated string mixing a vertical word ("auto"/"top"/"bottom") 288# with a side word ("left"/"right"); either may be omitted. The side picks which 289# flank of the line the rotated label sits on, VLINE_TEXT_OFFSET points away. 290_VLINE_EDGES: Final[tuple[str, ...]] = ("top", "bottom") 291_VLINE_SIDES: Final[tuple[str, ...]] = ("left", "right") 292 293 294def _parse_vline_loc(loc: object) -> tuple[str, str]: 295 """Split a loc string into (edge, side), order-free and each optional. 296 297 edge is "auto", "top" or "bottom"; side is "left" or "right". An omitted 298 edge falls back to "auto" (choose by the data) and an omitted side to 299 "right" (today's placement). Raises on an unrecognised or doubled word -- 300 each is a typo, not a choice. 301 """ 302 if not isinstance(loc, str): 303 raise TypeError(f"axvline 'loc' must be a string, got {type(loc)}") 304 edge: str | None = None 305 side: str | None = None 306 for tok in loc.split(): 307 if tok == "auto" or tok in _VLINE_EDGES: 308 if edge is not None: 309 raise ValueError(f"axvline 'loc' names two vertical positions, got {loc!r}") 310 edge = tok 311 elif tok in _VLINE_SIDES: 312 if side is not None: 313 raise ValueError(f"axvline 'loc' names two sides, got {loc!r}") 314 side = tok 315 else: 316 raise ValueError(f"axvline 'loc' word {tok!r} not recognised in {loc!r}") 317 return edge or "auto", side or "right" 318 319 320def _pop_vline_text(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]] | None: 321 """Remove the label keys from an axvline dict and return them, or None if unlabelled. 322 323 Raises on a bad loc, a non-dict text_kwargs, or label options given 324 without any text to place -- each of which is a typo, not a choice. 325 """ 326 text = item.pop("text", None) 327 loc = item.pop("loc", "auto") 328 text_kwargs = item.pop("text_kwargs", {}) 329 330 if text is None or not str(text).strip(): 331 if loc != "auto" or text_kwargs: 332 raise ValueError("axvline 'loc'/'text_kwargs' given without any 'text' to place") 333 return None 334 _parse_vline_loc(loc) # validate now; raises on a bad loc string 335 if not isinstance(text_kwargs, dict): 336 raise TypeError(f"axvline 'text_kwargs' must be a dict, got {type(text_kwargs)}") 337 return str(text), loc, text_kwargs 338 339 340def _data_y_extent(axes: Axes, x: float) -> tuple[float, float] | None: 341 """Return the (min, max) y of plotted data in a narrow x-band around x. 342 343 A rotated label occupies only a sliver of the x-axis, so a narrow band is 344 what the eye actually judges. Lines are matched by transform identity: 345 axhline/axvline use blended transforms, so this skips them and measures 346 only real data. Rectangles cover bar plots; axvspan/axhspan patches are 347 likewise blended and skipped. Returns None when nothing is measurable. 348 """ 349 left, right = axes.get_xlim() 350 half = abs(right - left) * VLINE_AUTO_BAND 351 low, high = x - half, x + half 352 found: list[float] = [] 353 354 for line in axes.get_lines(): 355 if not isinstance(line, Line2D) or line.get_transform() is not axes.transData: 356 continue 357 xdata = np.asarray(line.get_xdata(), dtype=float) 358 ydata = np.asarray(line.get_ydata(), dtype=float) 359 if xdata.size != ydata.size or xdata.size == 0: 360 continue 361 wanted = (xdata >= low) & (xdata <= high) & ~np.isnan(ydata) 362 if wanted.any(): 363 found.extend((float(ydata[wanted].min()), float(ydata[wanted].max()))) 364 365 for patch in axes.patches: 366 if not isinstance(patch, Rectangle) or patch.get_data_transform() is not axes.transData: 367 continue 368 x0, width = patch.get_x(), patch.get_width() 369 y0, height = patch.get_y(), patch.get_height() 370 if min(x0, x0 + width) > high or max(x0, x0 + width) < low: 371 continue 372 found.extend((min(y0, y0 + height), max(y0, y0 + height))) 373 374 if not found: 375 return None 376 return min(found), max(found) 377 378 379def _auto_vline_loc(axes: Axes, x: float) -> str: 380 """Pick the end of the axes with more room between the data and the limit.""" 381 extent = _data_y_extent(axes, x) 382 if extent is None: 383 return "top" # nothing measurable (empty band, or an unrecognised artist) 384 data_low, data_high = extent 385 bottom_lim, top_lim = axes.get_ylim() 386 if top_lim >= bottom_lim: 387 top_gap, bottom_gap = top_lim - data_high, data_low - bottom_lim 388 else: # inverted y-axis: values decrease going up the display 389 top_gap, bottom_gap = data_low - top_lim, bottom_lim - data_high 390 return "top" if top_gap >= bottom_gap else "bottom" 391 392 393def _annotate_vline(axes: Axes, item: dict[str, Any], line: Line2D, spec: tuple[str, str, dict]) -> None: 394 """Place a rotated text label just to one side of a vertical line. 395 396 The side (left/right of the line) comes from loc; the label is anchored 397 with x in data coordinates and y in axes coordinates, so it stays pinned 398 to the top/bottom of the plot regardless of any later change to the 399 y-limits. 400 """ 401 text, loc, text_kwargs = spec 402 x = item.get("x", 0) # matches the matplotlib default for axvline 403 edge, side = _parse_vline_loc(loc) 404 if edge == "auto": 405 edge = _auto_vline_loc(axes, float(x)) 406 y, valign = (1.0 - VLINE_TEXT_PAD, "top") if edge == "top" else (VLINE_TEXT_PAD, "bottom") 407 x_off, halign = (VLINE_TEXT_OFFSET, "left") if side == "right" else (-VLINE_TEXT_OFFSET, "right") 408 409 options: dict[str, Any] = { 410 "rotation": VLINE_TEXT_ROTATION, 411 "fontsize": VLINE_TEXT_FONTSIZE, 412 "color": line.get_color(), 413 "ha": halign, 414 "va": valign, 415 } 416 options.update(text_kwargs) 417 axes.annotate( 418 text, 419 xy=(x, y), 420 xycoords=blended_transform_factory(axes.transData, axes.transAxes), 421 xytext=(x_off, 0), 422 textcoords="offset points", 423 **options, 424 ) 425 426 427def _apply_splat(axes: Axes, method_name: str, value: _SplatValue) -> None: 428 """Apply a single splat kwarg, which may be a dict or sequence of dicts.""" 429 if value is None or value is False: 430 return 431 432 if value is True: # use the global default settings 433 value = get_setting(method_name) 434 435 # normalise to a list of dicts 436 if isinstance(value, dict): 437 value = [value] 438 439 if isinstance(value, Sequence): 440 method = getattr(axes, method_name) 441 for item in value: 442 if not isinstance(item, dict): 443 print(f"Warning: expected dict in {method_name} sequence, but got {type(item)}.") 444 continue 445 converted = _convert_period_coords(axes, method_name, item) 446 # _convert_period_coords always copies, so popping is safe here 447 spec = _pop_vline_text(converted) if method_name == "axvline" else None 448 artist = method(**converted) 449 if spec is not None and isinstance(artist, Line2D): 450 _annotate_vline(axes, converted, artist, spec) 451 else: 452 print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.") 453 454 455def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None: 456 """Set matplotlib elements dynamically using setting_name and splat.""" 457 for method_name in settings: 458 if method_name not in kwargs: 459 continue 460 461 if method_name == "legend": 462 legend_value = kwargs.get(method_name) 463 if isinstance(legend_value, (bool, dict, type(None))): 464 make_legend(axes, legend=legend_value) 465 else: 466 print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.") 467 continue 468 469 value = kwargs.get(method_name) 470 if value is None or isinstance(value, (bool, dict, Sequence)): 471 _apply_splat(axes, method_name, value) 472 else: 473 print(f"Warning: expected dict or sequence of dicts for {method_name}, but got {type(value)}.") 474 475 476def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 477 """Set figure size and apply chart annotations. 478 479 No-op when axes_only=True: the work here is all figure-level (resize, 480 corner text) and would stomp on other panels in a multi-axes figure. 481 """ 482 if kwargs.get("axes_only"): 483 return 484 fig = axes.figure 485 fig_size = kwargs.get("figsize", get_setting("figsize")) 486 if not isinstance(fig, SubFigure): 487 fig.set_size_inches(*fig_size) 488 489 annotations = { 490 "rfooter": (0.99, 0.001, "right", "bottom"), 491 "lfooter": (0.01, 0.001, "left", "bottom"), 492 "rheader": (0.99, 0.999, "right", "top"), 493 "lheader": (0.01, 0.999, "left", "top"), 494 } 495 496 for annotation in HEADER_FOOTER_KWARGS: 497 if annotation in kwargs: 498 x_pos, y_pos, h_align, v_align = annotations[annotation] 499 fig.text( 500 x_pos, 501 y_pos, 502 str(kwargs.get(annotation, "")), 503 ha=h_align, 504 va=v_align, 505 fontsize=FOOTNOTE_FONTSIZE, 506 fontstyle=FOOTNOTE_FONTSTYLE, 507 color=FOOTNOTE_COLOR, 508 ) 509 510 511def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 512 """Apply settings found in kwargs, after plotting the data.""" 513 apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs) 514 515 516def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 517 """Apply settings found in kwargs.""" 518 519 def check_kwargs(name: str) -> bool: 520 return name in kwargs and bool(kwargs.get(name)) 521 522 apply_value_kwargs(axes, VALUE_KWARGS, **kwargs) 523 apply_annotations(axes, **kwargs) 524 525 if check_kwargs("zero_y"): 526 bottom, top = axes.get_ylim() 527 adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT 528 if bottom > -adj: 529 axes.set_ylim(bottom=-adj) 530 if top < adj: 531 axes.set_ylim(top=adj) 532 533 if check_kwargs("y0"): 534 low, high = axes.get_ylim() 535 if low < 0 < high: 536 axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR) 537 538 if check_kwargs("x0"): 539 low, high = axes.get_xlim() 540 if low < 0 < high: 541 axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR) 542 543 if check_kwargs("axisbelow"): 544 axes.set_axisbelow(True) 545 546 547def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None: 548 """Save the figure to file.""" 549 saving = not kwargs.get("dont_save", False) # save by default 550 if not saving: 551 return 552 553 try: 554 chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir"))) 555 556 # Ensure directory exists 557 chart_dir.mkdir(parents=True, exist_ok=True) 558 559 suptitle = kwargs.get("suptitle", "") 560 title = kwargs.get("title", "") 561 pre_tag = kwargs.get("pre_tag", "") 562 tag = kwargs.get("tag", "") 563 name_override = kwargs.get("filename", "") 564 name_title = name_override or suptitle or title 565 file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME) 566 file_type = kwargs.get("file_type", get_setting("file_type")).lower() 567 dpi = kwargs.get("dpi", get_setting("dpi")) 568 569 # Construct filename components safely 570 filename_parts = [] 571 if pre_tag: 572 filename_parts.append(sanitize_filename(pre_tag)) 573 filename_parts.append(file_title) 574 if tag: 575 filename_parts.append(sanitize_filename(tag)) 576 577 # Join filename parts and add extension 578 filename = "-".join(filter(None, filename_parts)) 579 filepath = chart_dir / f"{filename}.{file_type}" 580 581 fig.savefig(filepath, dpi=dpi) 582 583 except ( 584 OSError, 585 PermissionError, 586 FileNotFoundError, 587 ValueError, 588 RuntimeError, 589 TypeError, 590 UnicodeError, 591 ) as e: 592 print(f"Error: Could not save plot to file: {e}") 593 594 595# - public functions for finalise_plot() 596 597 598def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 599 """Finalise and save plots to the file system. 600 601 The filename for the saved plot is constructed from the global 602 chart_dir, the plot's title, any specified tag text, and the 603 file_type for the plot. 604 605 Args: 606 axes: Axes - matplotlib axes object - required 607 kwargs: FinaliseKwargs 608 609 """ 610 # --- check the kwargs 611 report_kwargs(caller=ME, **kwargs) 612 validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs) 613 614 # --- sanity checks 615 if len(axes.get_children()) < 1: 616 print(f"Warning: {ME}() called with an empty axes, which was ignored.") 617 return 618 619 # --- remember axis-limits should we need to restore thems 620 xlim, ylim = axes.get_xlim(), axes.get_ylim() 621 622 # margins 623 axes.margins(DEFAULT_MARGIN) 624 axes.autoscale(tight=False) # This is problematic ... 625 626 apply_kwargs(axes, **kwargs) 627 628 # tight layout and save the figure 629 fig = axes.figure 630 axes_only = kwargs.get("axes_only", False) 631 if not axes_only and (suptitle := kwargs.get("suptitle")): 632 fig.suptitle(suptitle) 633 if kwargs.get("preserve_lims"): 634 # restore the original limits of the axes 635 axes.set_xlim(xlim) 636 axes.set_ylim(ylim) 637 if not axes_only and not isinstance(fig, SubFigure): 638 fig.tight_layout(pad=TIGHT_LAYOUT_PAD) 639 apply_late_kwargs(axes, **kwargs) 640 # axvspan/axvline in late_kwargs may have widened xlim beyond what 641 # set_labels() last saw; regenerate ticks from the updated view. 642 refresh_period_labels(axes) 643 # de-collide end-of-line annotations now that the layout/limits are final 644 resolve_annotation_collisions(axes) 645 legend = axes.get_legend() 646 if legend and kwargs.get("remove_legend", False): 647 legend.remove() 648 if not axes_only and not isinstance(fig, SubFigure): 649 save_to_file(fig, **kwargs) 650 651 # show the plot in Jupyter Lab 652 if not axes_only and kwargs.get("show"): 653 plt.show() 654 655 # And close - the figure this axes belongs to, not pyplot's current figure 656 if not axes_only and not kwargs.get("dont_close", False): 657 root = fig 658 while isinstance(root, SubFigure): 659 root = root.figure 660 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.
456def apply_splat_kwargs(axes: Axes, settings: tuple, **kwargs: Unpack[FinaliseKwargs]) -> None: 457 """Set matplotlib elements dynamically using setting_name and splat.""" 458 for method_name in settings: 459 if method_name not in kwargs: 460 continue 461 462 if method_name == "legend": 463 legend_value = kwargs.get(method_name) 464 if isinstance(legend_value, (bool, dict, type(None))): 465 make_legend(axes, legend=legend_value) 466 else: 467 print(f"Warning: expected bool, dict, or None for legend, but got {type(legend_value)}.") 468 continue 469 470 value = kwargs.get(method_name) 471 if value is None or isinstance(value, (bool, dict, Sequence)): 472 _apply_splat(axes, method_name, value) 473 else: 474 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.
477def apply_annotations(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 478 """Set figure size and apply chart annotations. 479 480 No-op when axes_only=True: the work here is all figure-level (resize, 481 corner text) and would stomp on other panels in a multi-axes figure. 482 """ 483 if kwargs.get("axes_only"): 484 return 485 fig = axes.figure 486 fig_size = kwargs.get("figsize", get_setting("figsize")) 487 if not isinstance(fig, SubFigure): 488 fig.set_size_inches(*fig_size) 489 490 annotations = { 491 "rfooter": (0.99, 0.001, "right", "bottom"), 492 "lfooter": (0.01, 0.001, "left", "bottom"), 493 "rheader": (0.99, 0.999, "right", "top"), 494 "lheader": (0.01, 0.999, "left", "top"), 495 } 496 497 for annotation in HEADER_FOOTER_KWARGS: 498 if annotation in kwargs: 499 x_pos, y_pos, h_align, v_align = annotations[annotation] 500 fig.text( 501 x_pos, 502 y_pos, 503 str(kwargs.get(annotation, "")), 504 ha=h_align, 505 va=v_align, 506 fontsize=FOOTNOTE_FONTSIZE, 507 fontstyle=FOOTNOTE_FONTSTYLE, 508 color=FOOTNOTE_COLOR, 509 )
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.
512def apply_late_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 513 """Apply settings found in kwargs, after plotting the data.""" 514 apply_splat_kwargs(axes, SPLAT_KWARGS, **kwargs)
Apply settings found in kwargs, after plotting the data.
517def apply_kwargs(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 518 """Apply settings found in kwargs.""" 519 520 def check_kwargs(name: str) -> bool: 521 return name in kwargs and bool(kwargs.get(name)) 522 523 apply_value_kwargs(axes, VALUE_KWARGS, **kwargs) 524 apply_annotations(axes, **kwargs) 525 526 if check_kwargs("zero_y"): 527 bottom, top = axes.get_ylim() 528 adj = (top - bottom) * ZERO_AXIS_ADJUSTMENT 529 if bottom > -adj: 530 axes.set_ylim(bottom=-adj) 531 if top < adj: 532 axes.set_ylim(top=adj) 533 534 if check_kwargs("y0"): 535 low, high = axes.get_ylim() 536 if low < 0 < high: 537 axes.axhline(y=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR) 538 539 if check_kwargs("x0"): 540 low, high = axes.get_xlim() 541 if low < 0 < high: 542 axes.axvline(x=0, lw=ZERO_LINE_WIDTH, c=ZERO_LINE_COLOR) 543 544 if check_kwargs("axisbelow"): 545 axes.set_axisbelow(True)
Apply settings found in kwargs.
548def save_to_file(fig: Figure, **kwargs: Unpack[FinaliseKwargs]) -> None: 549 """Save the figure to file.""" 550 saving = not kwargs.get("dont_save", False) # save by default 551 if not saving: 552 return 553 554 try: 555 chart_dir = Path(kwargs.get("chart_dir", get_setting("chart_dir"))) 556 557 # Ensure directory exists 558 chart_dir.mkdir(parents=True, exist_ok=True) 559 560 suptitle = kwargs.get("suptitle", "") 561 title = kwargs.get("title", "") 562 pre_tag = kwargs.get("pre_tag", "") 563 tag = kwargs.get("tag", "") 564 name_override = kwargs.get("filename", "") 565 name_title = name_override or suptitle or title 566 file_title = sanitize_filename(name_title or DEFAULT_FILE_TITLE_NAME) 567 file_type = kwargs.get("file_type", get_setting("file_type")).lower() 568 dpi = kwargs.get("dpi", get_setting("dpi")) 569 570 # Construct filename components safely 571 filename_parts = [] 572 if pre_tag: 573 filename_parts.append(sanitize_filename(pre_tag)) 574 filename_parts.append(file_title) 575 if tag: 576 filename_parts.append(sanitize_filename(tag)) 577 578 # Join filename parts and add extension 579 filename = "-".join(filter(None, filename_parts)) 580 filepath = chart_dir / f"{filename}.{file_type}" 581 582 fig.savefig(filepath, dpi=dpi) 583 584 except ( 585 OSError, 586 PermissionError, 587 FileNotFoundError, 588 ValueError, 589 RuntimeError, 590 TypeError, 591 UnicodeError, 592 ) as e: 593 print(f"Error: Could not save plot to file: {e}")
Save the figure to file.
599def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None: 600 """Finalise and save plots to the file system. 601 602 The filename for the saved plot is constructed from the global 603 chart_dir, the plot's title, any specified tag text, and the 604 file_type for the plot. 605 606 Args: 607 axes: Axes - matplotlib axes object - required 608 kwargs: FinaliseKwargs 609 610 """ 611 # --- check the kwargs 612 report_kwargs(caller=ME, **kwargs) 613 validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs) 614 615 # --- sanity checks 616 if len(axes.get_children()) < 1: 617 print(f"Warning: {ME}() called with an empty axes, which was ignored.") 618 return 619 620 # --- remember axis-limits should we need to restore thems 621 xlim, ylim = axes.get_xlim(), axes.get_ylim() 622 623 # margins 624 axes.margins(DEFAULT_MARGIN) 625 axes.autoscale(tight=False) # This is problematic ... 626 627 apply_kwargs(axes, **kwargs) 628 629 # tight layout and save the figure 630 fig = axes.figure 631 axes_only = kwargs.get("axes_only", False) 632 if not axes_only and (suptitle := kwargs.get("suptitle")): 633 fig.suptitle(suptitle) 634 if kwargs.get("preserve_lims"): 635 # restore the original limits of the axes 636 axes.set_xlim(xlim) 637 axes.set_ylim(ylim) 638 if not axes_only and not isinstance(fig, SubFigure): 639 fig.tight_layout(pad=TIGHT_LAYOUT_PAD) 640 apply_late_kwargs(axes, **kwargs) 641 # axvspan/axvline in late_kwargs may have widened xlim beyond what 642 # set_labels() last saw; regenerate ticks from the updated view. 643 refresh_period_labels(axes) 644 # de-collide end-of-line annotations now that the layout/limits are final 645 resolve_annotation_collisions(axes) 646 legend = axes.get_legend() 647 if legend and kwargs.get("remove_legend", False): 648 legend.remove() 649 if not axes_only and not isinstance(fig, SubFigure): 650 save_to_file(fig, **kwargs) 651 652 # show the plot in Jupyter Lab 653 if not axes_only and kwargs.get("show"): 654 plt.show() 655 656 # And close - the figure this axes belongs to, not pyplot's current figure 657 if not axes_only and not kwargs.get("dont_close", False): 658 root = fig 659 while isinstance(root, SubFigure): 660 root = root.figure 661 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