Source code for GUIBRUSHR.core.functions.figure_pickle

"""Helpers to persist matplotlib figures for interactive reload in the GUI.

Heavy plots are rendered in headless background subprocesses that only write
PNG/PDF to disk, so the live matplotlib ``Figure`` never reaches the GUI process.
To restore zoom/pan, each subprocess dumps the figure next to its PNG with
``dump_figure``; ``MyFigure`` then loads it and embeds an interactive canvas
(falling back to the static PNG when the dump is absent or unreadable).

Kept dependency-light (only ``pickle`` / ``os``) so subprocesses can import it
without pulling in the Tk backend.
"""

import os
import pickle

# Suffix appended to a PNG path to locate its pickled figure. Must stay in sync
# with the loader in GUIBRUSHR/GUI/WIDGET/MyFigure.py.
FIGPKL_SUFFIX = ".figpkl"

# Suffix used to store a list of figures (the pages of a multi-page PDF) next to
# that PDF. Must stay in sync with the loader in GUI/WIDGET/MyMultiFigure.py.
FIGS_SUFFIX = ".figs.pkl"




def figpkl_path(png_path) -> str:
    """Return the figpkl path that sits next to ``png_path``."""
    return str(png_path) + FIGPKL_SUFFIX





def figs_path(pdf_path) -> str:
    """Return the multi-figure pickle path for a PDF (``.pdf`` -> ``.figs.pkl``)."""
    text = str(pdf_path)
    if text.endswith(".pdf"):
        return text[: -len(".pdf")] + FIGS_SUFFIX
    return text + FIGS_SUFFIX





def dump_figure(fig, png_path) -> None:
    """Pickle a matplotlib Figure next to its PNG for interactive reload.

    Best-effort: any failure (an unpicklable figure, a read-only folder) is
    swallowed so it can never break the PNG/PDF generation that callers rely on.

    Args:
        fig: The matplotlib Figure to persist.
        png_path: Path of the PNG the figure was saved to; the pickle is written
            alongside it with the ``.figpkl`` suffix.
    """
    try:
        with open(figpkl_path(png_path), "wb") as handle:
            pickle.dump(fig, handle)
    except Exception:
        pass





def load_figure(png_path):
    """Load a pickled figure for ``png_path`` if present, else return ``None``.

    Best-effort: returns ``None`` on a missing file or any unpickling error so
    the caller can fall back to the static image.
    """
    path = figpkl_path(png_path)
    if not os.path.exists(path):
        return None
    try:
        with open(path, "rb") as handle:
            return pickle.load(handle)
    except Exception:
        return None





def save_figure_blobs(blobs, pdf_path) -> None:
    """Persist a list of already-pickled figure byte-blobs next to a PDF.

    The on-disk ``.figs.pkl`` format is a list of ``pickle.dumps(figure)`` blobs.
    Storing bytes (rather than live figures) is deliberate: unpickling a figure
    re-registers it with pyplot's figure manager, so a retained live copy would
    be wiped by a later ``plt.clf()`` in the page-producing code. Bytes captured
    at draw time are an immutable snapshot, immune to that.

    Best-effort: any failure is swallowed so it can never break the PDF.
    """
    try:
        with open(figs_path(pdf_path), "wb") as handle:
            pickle.dump(list(blobs), handle)
    except Exception:
        pass





def dump_figures(figures, pdf_path) -> None:
    """Serialize each figure to bytes immediately, then persist the blob list.

    Convenience for callers holding stable live figures. Each figure is pickled
    to bytes right away (see ``save_figure_blobs`` for why bytes).
    """
    blobs = []
    for figure in figures:
        try:
            blobs.append(pickle.dumps(figure))
        except Exception:
            pass
    save_figure_blobs(blobs, pdf_path)





def load_figures(pdf_path):
    """Load the figures for ``pdf_path`` (unpickling each blob), else ``None``.

    Best-effort: returns ``None`` on a missing file or any unpickling error so
    the caller can fall back to the static PDF viewer.
    """
    path = figs_path(pdf_path)
    if not os.path.exists(path):
        return None
    try:
        with open(path, "rb") as handle:
            blobs = pickle.load(handle)
    except Exception:
        return None
    figures = []
    for blob in blobs or []:
        try:
            figures.append(pickle.loads(blob) if isinstance(blob, (bytes, bytearray)) else blob)
        except Exception:
            pass
    return figures or None