|
| 1 | +"""sovereign_debt_py.plotting.core |
| 2 | +
|
| 3 | +Internal helpers: validation, styling, figure creation, and the |
| 4 | +fig_to_png_bytes rendering utility. |
| 5 | +""" |
| 6 | + |
| 7 | +from __future__ import annotations |
| 8 | + |
| 9 | +import io |
| 10 | +import datetime |
| 11 | +from typing import Sequence, Any |
| 12 | + |
| 13 | +import numpy as np |
| 14 | +import matplotlib |
| 15 | +import matplotlib.pyplot as plt |
| 16 | +import matplotlib.ticker as mticker |
| 17 | +from matplotlib.figure import Figure |
| 18 | +from matplotlib.axes import Axes |
| 19 | + |
| 20 | + |
| 21 | +# --------------------------------------------------------------------------- |
| 22 | +# Input validation |
| 23 | +# --------------------------------------------------------------------------- |
| 24 | + |
| 25 | +def to_1d_array(x: Any) -> np.ndarray: |
| 26 | + """Convert list, tuple, numpy array, or pandas Series to a 1-D ndarray. |
| 27 | +
|
| 28 | + Raises |
| 29 | + ------ |
| 30 | + ValueError |
| 31 | + If *x* cannot be converted or is not 1-D after conversion. |
| 32 | + """ |
| 33 | + try: |
| 34 | + # pandas Series / Index |
| 35 | + arr = np.asarray(x, dtype=float) |
| 36 | + except (TypeError, ValueError) as exc: |
| 37 | + raise ValueError( |
| 38 | + f"Cannot convert input to a numeric 1-D array: {exc}" |
| 39 | + ) from exc |
| 40 | + |
| 41 | + if arr.ndim != 1: |
| 42 | + raise ValueError( |
| 43 | + f"Expected a 1-D array-like, got shape {arr.shape}." |
| 44 | + ) |
| 45 | + return arr |
| 46 | + |
| 47 | + |
| 48 | +def validate_same_length(*arrays: np.ndarray) -> None: |
| 49 | + """Raise ValueError if any two arrays differ in length.""" |
| 50 | + lengths = [len(a) for a in arrays] |
| 51 | + if len(set(lengths)) > 1: |
| 52 | + raise ValueError( |
| 53 | + f"All inputs must have the same length; got lengths: {lengths}." |
| 54 | + ) |
| 55 | + |
| 56 | + |
| 57 | +def coerce_dates(dates: Sequence[Any]) -> list[datetime.datetime]: |
| 58 | + """Convert a sequence of date-like objects to :class:`datetime.datetime`. |
| 59 | +
|
| 60 | + Accepts ``datetime.date``, ``datetime.datetime``, ISO-format strings, |
| 61 | + and numpy / pandas timestamp-like objects. |
| 62 | +
|
| 63 | + Raises |
| 64 | + ------ |
| 65 | + ValueError |
| 66 | + If any element cannot be parsed. |
| 67 | + """ |
| 68 | + result: list[datetime.datetime] = [] |
| 69 | + for item in dates: |
| 70 | + if isinstance(item, datetime.datetime): |
| 71 | + result.append(item) |
| 72 | + elif isinstance(item, datetime.date): |
| 73 | + result.append(datetime.datetime(item.year, item.month, item.day)) |
| 74 | + elif isinstance(item, str): |
| 75 | + try: |
| 76 | + result.append(datetime.datetime.fromisoformat(item)) |
| 77 | + except ValueError as exc: |
| 78 | + raise ValueError( |
| 79 | + f"Cannot parse date string {item!r}: {exc}" |
| 80 | + ) from exc |
| 81 | + else: |
| 82 | + # numpy datetime64, pandas Timestamp, etc. |
| 83 | + try: |
| 84 | + import pandas as pd # optional |
| 85 | + ts = pd.Timestamp(item) |
| 86 | + result.append(ts.to_pydatetime()) |
| 87 | + except Exception: |
| 88 | + # Last resort: try numpy datetime64 → Python datetime |
| 89 | + try: |
| 90 | + ts_ms = np.datetime64(item, "ms") |
| 91 | + epoch = np.datetime64(0, "ms") |
| 92 | + ms = int((ts_ms - epoch) / np.timedelta64(1, "ms")) |
| 93 | + result.append( |
| 94 | + datetime.datetime(1970, 1, 1) |
| 95 | + + datetime.timedelta(milliseconds=ms) |
| 96 | + ) |
| 97 | + except Exception as exc2: |
| 98 | + raise ValueError( |
| 99 | + f"Cannot coerce {item!r} to datetime: {exc2}" |
| 100 | + ) from exc2 |
| 101 | + return result |
| 102 | + |
| 103 | + |
| 104 | +# --------------------------------------------------------------------------- |
| 105 | +# Styling |
| 106 | +# --------------------------------------------------------------------------- |
| 107 | + |
| 108 | +def apply_sdpy_style(ax: Axes, *, grid: bool = True) -> None: |
| 109 | + """Apply consistent sovereign_debt_py chart styling to *ax*. |
| 110 | +
|
| 111 | + Adjusts spines, font sizes, line widths, and optionally adds a grid. |
| 112 | + Does **not** call ``plt.style.use`` globally. |
| 113 | + """ |
| 114 | + ax.spines["top"].set_visible(False) |
| 115 | + ax.spines["right"].set_visible(False) |
| 116 | + ax.spines["left"].set_linewidth(0.8) |
| 117 | + ax.spines["bottom"].set_linewidth(0.8) |
| 118 | + |
| 119 | + ax.tick_params(axis="both", labelsize=9, length=4, width=0.8) |
| 120 | + ax.xaxis.label.set_size(10) |
| 121 | + ax.yaxis.label.set_size(10) |
| 122 | + if ax.get_title(): |
| 123 | + ax.title.set_size(11) |
| 124 | + ax.title.set_fontweight("bold") |
| 125 | + |
| 126 | + if grid: |
| 127 | + ax.grid(True, linestyle="--", linewidth=0.5, alpha=0.6) |
| 128 | + ax.set_axisbelow(True) |
| 129 | + else: |
| 130 | + ax.grid(False) |
| 131 | + |
| 132 | + |
| 133 | +# --------------------------------------------------------------------------- |
| 134 | +# Figure factory |
| 135 | +# --------------------------------------------------------------------------- |
| 136 | + |
| 137 | +def _make_fig_ax( |
| 138 | + fig: Figure | None, |
| 139 | + ax: Axes | None, |
| 140 | +) -> tuple[Figure, Axes]: |
| 141 | + """Return *(fig, ax)*, creating new ones if not supplied.""" |
| 142 | + if fig is None and ax is None: |
| 143 | + fig, ax = plt.subplots() |
| 144 | + elif fig is None: |
| 145 | + fig = ax.get_figure() |
| 146 | + elif ax is None: |
| 147 | + ax = fig.gca() |
| 148 | + return fig, ax # type: ignore[return-value] |
| 149 | + |
| 150 | + |
| 151 | +# --------------------------------------------------------------------------- |
| 152 | +# Rendering helper |
| 153 | +# --------------------------------------------------------------------------- |
| 154 | + |
| 155 | +def fig_to_png_bytes( |
| 156 | + fig: Figure, |
| 157 | + *, |
| 158 | + width_px: int = 800, |
| 159 | + height_px: int = 450, |
| 160 | + dpi: int = 120, |
| 161 | + tight: bool = True, |
| 162 | + close: bool = False, |
| 163 | +) -> bytes: |
| 164 | + """Render *fig* to PNG and return the raw bytes. |
| 165 | +
|
| 166 | + The figure's size is set (in inches) based on ``width_px / dpi`` and |
| 167 | + ``height_px / dpi`` before saving, so the output is deterministic. |
| 168 | +
|
| 169 | + Parameters |
| 170 | + ---------- |
| 171 | + fig: |
| 172 | + A Matplotlib :class:`~matplotlib.figure.Figure`. |
| 173 | + width_px: |
| 174 | + Output width in pixels. |
| 175 | + height_px: |
| 176 | + Output height in pixels. |
| 177 | + dpi: |
| 178 | + Dots per inch for the PNG encoder. |
| 179 | + tight: |
| 180 | + If *True*, call ``tight_layout()`` before saving. |
| 181 | + close: |
| 182 | + If *True*, close *fig* after saving (useful for embedding workflows). |
| 183 | +
|
| 184 | + Returns |
| 185 | + ------- |
| 186 | + bytes |
| 187 | + Raw PNG bytes (starts with ``b'\\x89PNG'``). |
| 188 | + """ |
| 189 | + fig.set_size_inches(width_px / dpi, height_px / dpi) |
| 190 | + if tight: |
| 191 | + fig.tight_layout() |
| 192 | + |
| 193 | + buf = io.BytesIO() |
| 194 | + # Use the Agg backend for headless rendering without touching the global |
| 195 | + # backend. canvas.print_figure works regardless of the current backend. |
| 196 | + agg_canvas = matplotlib.backends.backend_agg.FigureCanvasAgg(fig) |
| 197 | + agg_canvas.print_figure(buf, format="png", dpi=dpi) |
| 198 | + png_bytes = buf.getvalue() |
| 199 | + |
| 200 | + if close: |
| 201 | + plt.close(fig) |
| 202 | + |
| 203 | + return png_bytes |
0 commit comments