How to overwrite `savefig` method so that it can unpack a tuple

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')

Solution

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')