Ejemplo n.º 1
0
def loadFigureFromFile(filename: str,
                       figure: Figure = None,
                       offset: list = None,
                       dpi: int = None,
                       cache: bool = False):
    """
    Add contents to the current figure from the file defined by filename. It can be either a python script defining
    a figure, an image (filename or directly the numpy array), or an svg file.

    See also :ref:`composing`.

    Parameters
    ----------
    filename : str
        The file to load. Can point to a python script file, an image file or an svg file.
    figure : Figure, optional
        The figure where to add the loaded file. Defaults to the current figure.
    offset : list, optional
        The offset where to import the file. The first two parts define the x and y position and the third part defines
        the units to use. Default is "%", a percentage of the current figure size. It can also be "cm" or "in".
    cache : bool, optional
        Whether to try to cache the figure generated from the file. Only for python files. This option is experimental
        and may not be stable.
    """
    from matplotlib import rcParams
    from pylustrator import changeFigureSize
    import pylustrator

    # change to the directory of the filename (to execute the code relative to this directory)
    dirname, filename = os.path.split(filename)
    dirname = os.path.abspath(dirname)
    with changeFolder(dirname):
        if dirname:
            os.chdir(dirname)

        # defaults to the current figure
        if figure is None:
            figure = plt.gcf()

        class noShow:
            """
            An environment that prevents the script from calling the plt.show function
            """
            def __enter__(self):
                # store the show function
                self.show = plt.show
                self.dragger = pylustrator.start

                # define an empty function
                def empty(*args, **kwargs):
                    pass

                # set the show function to the empty function
                plt.show = empty
                pylustrator.start = empty

            def __exit__(self, type, value, traceback):
                # restore the old show function
                plt.show = self.show
                pylustrator.start = self.dragger

        class noNewFigures:
            """
            An environment that prevents the script from creating new figures in the figure manager
            """
            def __enter__(self):
                fig = plt.gcf()
                self.fig = plt.figure
                figsize = rcParams['figure.figsize']
                fig.set_size_inches(figsize[0], figsize[1])

                def figure(num=None, figsize=None, *args, **kwargs):
                    fig = plt.gcf()
                    if figsize is not None:
                        fig.set_size_inches(figsize[0],
                                            figsize[1],
                                            forward=True)
                    return fig

                plt.figure = figure

            def __exit__(self, type, value, traceback):
                from matplotlib.figure import Figure
                from matplotlib.transforms import TransformedBbox, Affine2D
                plt.figure = self.fig

        # get the size of the old figure
        w1, h1 = figure.get_size_inches()
        axes1 = removeContentFromFigure(figure)
        if len(axes1) == 0:
            w1 = 0
            h1 = 0

        # try to load the filename as an image
        try:
            im = plt.imread(filename)
        except OSError:
            im = None

        # if it is an image, just display the image
        if im is not None:
            im = plt.imread(filename)
            imShowFullFigure(im, os.path.split(filename)[1], figure, dpi=dpi)
        # if the image is a numpy array, just display the array
        elif isinstance(filename, np.ndarray):
            im = filename
            imShowFullFigure(im, str(im.shape), figure, dpi)
        # if it is a svg file, display the svg file
        elif filename.endswith(".svg"):
            svgread(filename)
        # if not, it should be a python script
        else:
            filename = os.path.abspath(filename)
            cache_filename = filename + ".cache.pkl"

            with noNewFigures():
                # prevent the script we want to load from calling show
                with noShow():
                    import pickle
                    if cache and os.path.exists(
                            cache_filename) and os.path.getmtime(
                                cache_filename) > os.path.getmtime(filename):
                        print("loading from cached file", cache_filename)
                        fig2 = pickle.load(open(cache_filename, "rb"))
                        w, h = fig2.get_size_inches()
                        figure.set_size_inches(w, h)

                        str(figure
                            )  # important! (for some reason I don't know)
                        for ax in fig2.axes:
                            fig2.delaxes(ax)
                            figure._axstack.add(figure._make_key(ax), ax)
                            figure.bbox._parents.update(fig2.bbox._parents)
                            figure.dpi_scale_trans._parents.update(
                                fig2.dpi_scale_trans._parents)
                            replace_all_refs(fig2.bbox, figure.bbox)
                            replace_all_refs(fig2.dpi_scale_trans,
                                             figure.dpi_scale_trans)
                            replace_all_refs(fig2, figure)
                    else:
                        # execute the file
                        exec(
                            compile(
                                open(filename, "rb").read(), filename, 'exec'),
                            globals())
                        if cache is True:
                            c = figure.canvas
                            figure.canvas = None
                            figure.bbox.pylustrator = True
                            figure.dpi_scale_trans.pylustrator = True
                            pickle.dump(figure, open(cache_filename, 'wb'))

                            figure.canvas = c

        # get the size of the new figure
        w2, h2 = figure.get_size_inches()
        if offset is not None:
            if len(offset) == 2 or offset[2] == "%":
                w2 += w1 * offset[0]
                h2 += h1 * offset[1]
            elif offset[2] == "in":
                w2 += offset[0]
                h2 += offset[1]
            elif offset[2] == "cm":
                w2 += offset[0] / 2.54
                h2 += offset[1] / 2.54
            changeFigureSize(w2,
                             h2,
                             cut_from_top=True,
                             cut_from_left=True,
                             fig=figure)
        w = max(w1, w2)
        h = max(h1, h2)
        changeFigureSize(w, h, fig=figure)
        if len(axes1):
            axes2 = removeContentFromFigure(figure)
            changeFigureSize(w1, h1, fig=figure)
            addContentToFigure(figure, axes1)

            changeFigureSize(w, h, fig=figure)
            addContentToFigure(figure, axes2)