Is there a way to change the stylesheet when saving a figure in matplotlib?

Question

Is it possible to temporarily change the style sheet of a figure when saving it, without having to replot it (twice)? I know that you can pass the facecolor and edgecolor of the figure as parameters to savefig, but I don't understand why it would be limited to only these two parameters. You don't seem to have any influence on the rcParams of the axes instance(s). Matplotlib's style sheet context manager only affects the style sheet of the figure. Here is an MWE to play around with:

import matplotlib.pyplot as plt


fig, ax = plt.subplots()
ax.plot((1, 2), (0, 0), '--')

with plt.style.context('dark_background'):
    # if you replot it, it works, however, this messses up the style outside
    # this context as well
    # fig.clf()
    # ax = fig.add_subplot(111)
    # ax.plot((1, 2), (0, 0), '--')
    fig.savefig('test.png')

# you would need to replot everything again for the style to be correct again
# fig.clf()
# ax = fig.add_subplot(111)
# ax.plot((1, 2), (0, 0), '--')
fig.show()

Answer 1

Ok, so after a whole lot of research ad trial and error, I have finally found a way that works, although complex. The basic idea is to write your own ('smart') context manager, that understands what you want to do on an abstract level, and then makes the appropriate changes. E.g. if you want to change the axes edgecolor of a figure it knows which color changes are associated with it.

In the case of changing the edgecolor, there are quite a few obstacles that you need to be aware of, since changing the (edge-)color of the associated artists unfortunately works differently for different artists. Here are some of the caveats that I came across, the list is of course by no means comprehensive, since I am sore for more complex figures, there are also other things that you need to take into consideration:

• While the facecolor of an axes instance can simply be changed by using axes.set_facecolor(), for some reason, there is no equivalent method for the edgecolor. So, you need to do it by hand and collect all the spine, label and tick instances and change their colors individually.
• While the spine artist accepts the color, facecolor, as well as edgecolor keywords, and even has a set_color() method, it does not come with a get_color() method. Instead, you need to call get_edgecolor().
• If on top of that, you also happen to have a colorbar incorporated in your plot, you have to find yet another approach, because if you want to change the spine color of the colorbar, you're not doing that by changing the color of the spines, but by changing the color of the colorbar.outline. Which for some reason, doesn't seem to be documented in the official documentation. But if you want to change an aspect of the colorbar, you somehow need to retrieve it from the figure instance first.

I'm sure there are good reasons why the things work the way they do, but from the simple idea of wanting to change the edgecolor, it seems unnecessarily arbitrary.

Below is the code of the context manager applied to a simple example figure. I am sure there are many ways of how to improve this approach though.

import matplotlib.pyplot as plt
from matplotlib.spines import Spine
from matplotlib.patches import Polygon


class StyleChange:
    def __init__(self, color):
        self._color = color
        self._original_colors = {}

    def apply(self, fig):
        pass

    def revert(self):
        pass


class AxesEdgeColorChange(StyleChange):
    def apply(self, fig):
        for ax in fig.axes:
            self._original_colors[ax] = {}
            ticks = [*ax.get_xticklines(), *ax.get_yticklines()]
            mticks = [
                *ax.get_xaxis().get_minor_ticks(),
                *ax.get_yaxis().get_minor_ticks(),
            ]
            labels = [*ax.get_xticklabels(), *ax.get_yticklabels()]
            spines = ax.spines.values()
            cbars = [
                im.colorbar.outline for im in ax.images
                if (im is not None and im.colorbar is not None)
            ]
            for artist in [*ticks, *mticks, *labels, *spines, *cbars]:
                self._original_colors[ax][artist] = self._get_color(artist)
                self._set_color(artist)

    def revert(self):
        for axes_artists in self._original_colors.values():
            for artist, color in axes_artists.items():
                self._set_color(artist, color)

    @staticmethod
    def _get_color(artist):
        if isinstance(artist, (Spine, Polygon)):
            return artist.get_edgecolor()

        return artist.get_color()

    def _set_color(self, artist, color=None):
        if color is None:
            color = self._color
        if isinstance(artist, Polygon):
            artist.set_edgecolor(color)
        else:
            artist.set_color(color)


class AxesFaceColorChange(StyleChange):
    def apply(self, fig):
        for ax in fig.axes:
            self._original_colors[ax] = ax.get_facecolor()
            ax.set_facecolor(self._color)

    def revert(self):
        for ax, color in self._original_colors.items():
            ax.set_facecolor(color)


class StyleChangeManager:
    def __init__(
            self, fig: plt.Figure, *style_changes: StyleChange,
            redraw: bool = False
    ):
        self._fig = fig
        self._style_changes = style_changes
        self._redraw = redraw

    def __enter__(self):
        for change in self._style_changes:
            change.apply(self._fig)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        for change in self._style_changes:
            change.revert()
        if self._redraw:
            self._fig.canvas.draw_idle()


def example():
    image = np.random.uniform(0, 1, (100, 100))
    fig, ax = plt.subplots()
    image_artist = ax.imshow(image)
    fig.colorbar(image_artist)
    ec_change = AxesEdgeColorChange('red')
    with StyleChangeManager(fig, ec_change):
        fig.show()

    fig.show()


if __name__ == '__main__':
    import numpy as np

    example()