I am working on a minimal extension to a piece of code.
I have a make_fig
function that used to produce just one figure, which I would call as fig
in the lots of other fuctions then save it as fig.savefig
.
In the extension, make_fig
now returns a tuple of figures. So, to save them all, I would now need need something like:
fig = make_fig for f in fig: f.savefig
I was hoping for a more elegant solution, one that does not require a for loop to be added anywhere make_fig
appears.
Can I somehow modify the a.savefig
method so that it does its normal function if a
is a matplotlib.pytplot
instance, and does the for loop thing above if it’s a tuple?
MWE below.
d = 1
is the “old” code, d=2
is the extension I want to add.
import matplotlib.pyplot as plt def make_fig(d): if d==1: fig, ax = plt.subplots(1) elif d==2: fig1, ax = plt.subplots(1) fig2, ax = plt.subplots(1) fig = (fig1, fig2) else: raise Exception('error') return fig d=2 fig = make_fig(d) fig.savefig('hello.png')
Advertisement
Answer
Simply implement your own savefig
function that can handle both cases.
from collections.abc import Iterable def savefig(filename, fig): if not isinstance(fig, Iterable): fig = (fig,) for i, f in enumerate(fig): f.savefig(f'{filename}_{i}.jpg') fig = make_fig(1) savefig('test1', fig) fig = make_fig(2) savefig('test2', fig)
After execution, we have test1_0.jpg
, test2_0.jpg
and test2_1.jpg
.
As an alternative to the if
check, you can use the EAFP approach:
def savefig(filename, fig): try: fig.savefig(filename) except AttributeError: for i, f in enumerate(fig): f.savefig(f'{filename}_{i}.jpg')