Source code for gofigr.backends.matplotlib

"""\
Copyright (c) 2023, Flagstaff Solutions, LLC
All rights reserved.

"""
import inspect
import io

import matplotlib
import matplotlib.pyplot as plt

from gofigr.backends import GoFigrBackend, get_all_function_arguments




class SuppressPltWarnings:
    """Temporarily suppresses matplotlib warnings."""

    def __init__(self):
        self.plt_log = None
        self.log_level = None

    def __enter__(self):
        # Save initial log level
        self.plt_log = getattr(plt, "_log")
        self.log_level = getattr(self.plt_log, "level") if self.plt_log is not None else None
        plt.set_loglevel("error")

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Restore log level
        if self.plt_log is not None and self.log_level is not None:
            self.plt_log.setLevel(self.log_level)





class MatplotlibBackend(GoFigrBackend):
    """\
    MatplotLib backend for GoFigr.

    """


    def get_backend_name(self):
        return "matplotlib"




    def is_compatible(self, fig):
        return isinstance(fig, matplotlib.figure.Figure)




    def is_interactive(self, fig):
        return False




    def find_figures(self, shell, data):
        frames = inspect.stack()
        # Walk through the stack in *reverse* order (from top to bottom), to find the first call
        # in case display() was called recursively
        for f in reversed(frames):
            if f.function == "display" and ("IPython" in f.filename or 'matplot' in f.filename):
                for arg_value in get_all_function_arguments(f):
                    if self.is_compatible(arg_value):
                        yield arg_value

                break




    def get_default_figure(self, silent=False):
        return plt.gcf()




    @staticmethod
    def title_to_string(title):
        """Extracts the title as a string from a title-like object (e.g. Text)"""
        if title is None:
            return None
        elif isinstance(title, matplotlib.text.Text):
            return title.get_text()
        elif isinstance(title, str):
            return title
        else:
            return None




    def get_title(self, fig):
        suptitle = MatplotlibBackend.title_to_string(getattr(fig, "_suptitle", ""))
        # get_title() only reads the requested slot, and a figure titled with
        # loc="left"/"right" has an empty center slot — check all three so
        # such figures don't publish as 'Anonymous Figure'.
        title = None
        if len(fig.axes) > 0:
            for loc in ("center", "left", "right"):
                title = MatplotlibBackend.title_to_string(fig.axes[0].get_title(loc=loc))
                if title is not None and title.strip() != "":
                    break
        if suptitle is not None and suptitle.strip() != "":
            return suptitle
        elif title is not None and title.strip() != "":
            return title
        else:
            return None




    def figure_to_bytes(self, fig, fmt, params):
        with SuppressPltWarnings():
            params = dict(params)
            if "dpi" not in params:
                # Resolve the dpi explicitly instead of letting savefig read
                # ambient state: fig.dpi can be transiently scaled when we
                # observe the figure mid-render (e.g. IPython's retina
                # display doubles it inside print_figure), which otherwise
                # inflates the published image. _original_dpi is the stable
                # value matplotlib itself uses for savefig.dpi='figure'.
                dpi = matplotlib.rcParams["savefig.dpi"]
                if dpi == "figure":
                    dpi = getattr(fig, "_original_dpi", fig.dpi)
                params["dpi"] = dpi
            bio = io.BytesIO()
            fig.savefig(bio, format=fmt, **params)
            bio.seek(0)
            return bio.read()




    def close(self, fig):
        plt.close(fig)




    def get_supported_image_formats(self):
        return ["png", "eps", "pdf", "jpeg", "jpg", "svg", "tiff"]