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

Question

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

DeepSpace · Accepted Answer · 2021-01-17T22:09:02.290

1

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

edited Jan 17 '21 at 22:09

answered Jan 17 '21 at 21:41

DeepSpace

78,697
11
109
154

thanks. inside your `custom_savefig`, how would I refer to `fig`? – SuperCiocia Jan 17 '21 at 21:44
@SuperCiocia with `self`, I updated my answer – DeepSpace Jan 17 '21 at 21:45
But this only works if `fig` is a Figure, so that it has `savefig` as a method. What if my `fig = (fig1, fig2)`, where `fig1` and `fig2` are Figures? – SuperCiocia Jan 17 '21 at 21:51
@SuperCiocia See updated answer, a simpler and much safer solution – DeepSpace Jan 17 '21 at 22:02
@SuperCiocia `from collections.abc import Iterable` – DeepSpace Jan 17 '21 at 22:06
Hi @DeepSpace, I think this issue almost similar to another issue https://stackoverflow.com/q/68348996/6446053 . Appreciate if you can drop by and share some insight – mpx Jul 12 '21 at 14:49

score 1 · Answer 2 · answered Jan 17 '21 at 22:09

In order to do that, rather than returning a tuple, you can return your own object, which has the tuple as one of the fields.

class MultipleFigures:

    def __init__(self, figures):
        self.figures = figures

    def savefig(self):
        for fig in self.figures:
            fig.savefig()

The make fig function can then return either a matplotlib.pytplot instance or a MultipleFigures instance.

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

2 Answers2