backyard/library/plots.py

import matplotlib.pyplot as plt


class Plot:
    defaults = dict(figsize=(6, 4), dpi=128)

    def __init__(self, show=True, **kwargs):
        self.kwargs = self.defaults.copy()
        self.kwargs.update(kwargs)
        self.show = show

    def __enter__(self):
        self.fig = plt.figure(**self.kwargs)
        self.ax = self.fig.gca()
        self.fig.patch.set_facecolor("white")  # hacky
        return self.fig, self.ax

    def __exit__(self, exc_type, exc_val, exc_tb):
        from sys import stdout, stderr
        stdout.flush()  # doesn't work?
        stderr.flush()  # doesn't work?
        if self.show == "terminal":
            from io import BytesIO
            from autopsy import termshow
            buf = BytesIO()
            self.fig.savefig(buf, dpi=self.kwargs.get("dpi", 96), format="png")
            termshow(buf)
        elif self.show:
            plt.show()
        plt.close(self.fig)


class SquarePlot(Plot):  # TODO: delete me or fix me; not terribly useful.
    defaults = dict(figsize=(7, 7), dpi=128)

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def __enter__(self):
        ret = super().__enter__()
        self.ax.set_aspect("equal", "box")
        self.ax.grid(True, "major", alpha=0.8)
        self.ax.grid(True, "minor", alpha=0.4)
        return ret


class CleanPlot(Plot):
    def __enter__(self):
        super().__enter__()
        self.ax.set_axis_off()
        self.ax.set_position([0, 0, 1, 1])
        return self.fig, self.ax


class CleanPlotUnity(CleanPlot):
    def __init__(self, size=1, dim=None, **kwargs):
        if "figsize" not in kwargs and size is not None:
            kwargs["figsize"] = (size, size)
        if "dpi" not in kwargs and dim is not None:
            kwargs["dpi"] = dim
        super().__init__(**kwargs)

    def __enter__(self):
        super().__enter__()
        self.ax.set_xlim(0, 1)
        self.ax.set_ylim(0, 1)
        return self.fig, self.ax


class SquareCenterPlot(Plot):
    defaults = dict(figsize=(7, 7), dpi=128)

    def __init__(self, size=5, subdivisions=4, **kwargs):
        self.size = size
        self.subdivisions = subdivisions
        super().__init__(**kwargs)

    def __enter__(self):
        ret = super().__enter__()
        self.ax.set_aspect("equal", "box")
        self.ax.set_xlim(-self.size, self.size)
        self.ax.set_ylim(-self.size, self.size)
        self.ax.spines["left"].set_position("zero")
        self.ax.spines["bottom"].set_position("zero")
        self.ax.spines["right"].set_color("none")
        self.ax.spines["top"].set_color("none")
        self.ax.xaxis.set_ticks_position("bottom")
        self.ax.yaxis.set_ticks_position("left")
        p = self.subdivisions
        ticks = [x for x in range(-self.size * p, self.size * p + 1) if x != 0]
        major_ticks = [x / p for x in ticks if x % p == 0]
        minor_ticks = [x / p for x in ticks if x % p != 0]
        self.ax.set_xticks(major_ticks, minor=False)
        self.ax.set_yticks(major_ticks, minor=False)
        self.ax.set_xticks(minor_ticks, minor=True)
        self.ax.set_yticks(minor_ticks, minor=True)
        self.ax.grid(True, "major", alpha=0.8)
        self.ax.grid(True, "minor", alpha=0.4)
        return ret


class GridPlot:
    defaults = dict(figsize=(9, 6), dpi=128)

    def __init__(self, rows, cols, show=True, **kwargs):
        self.rows = int(rows)
        self.cols = int(cols)
        self.kwargs = self.defaults.copy()
        self.kwargs.update(kwargs)
        self.show = bool(show)

    def __enter__(self):
        self.fig, self.axes = plt.subplots(self.rows, self.cols, **self.kwargs)
        self.fig.patch.set_facecolor("white")
        return self.fig, self.axes

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.fig.tight_layout()
        if self.show:
            plt.show()
        plt.close(self.fig)