如何覆盖`savefig`方法以便它可以解包元组

如何解决如何覆盖`savefig`方法以便它可以解包元组

我正在对一段代码进行最小扩展。

我有一个 make_fig 函数，它曾经只生成一个图形，我会在许多其他函数中将其称为 fig，然后将其保存为 fig.savefig。

在扩展中，make_fig 现在返回一个元组数字。所以，为了保存它们，我现在需要这样的东西：

fig = make_fig

for f in fig:
   f.savefig

我希望有一种更优雅的解决方案，不需要在出现 make_fig 的任何地方添加 for 循环。

如果 a.savefig 是一个 a 实例，我可以以某种方式修改 matplotlib.pytplot 方法，使其执行其正常功能，并且如果它是一个元组，则执行上面的 for 循环吗？

>

下面的 MWE。

d = 1 是“旧”代码，d=2 是我要添加的扩展。

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

解决方法

只需实现您自己的 savefig 函数即可处理这两种情况。

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)

执行后，我们有 test1_0.jpg、test2_0.jpg 和 test2_1.jpg。

作为 if 检查的替代方法，您可以使用 EAFP 方法：

def savefig(filename,fig):
    try:
        fig.savefig(filename)
    except AttributeError:
        for i,f in enumerate(fig):
            f.savefig(f'{filename}_{i}.jpg')

,

为了做到这一点，您可以返回自己的对象，而不是返回元组，该对象将元组作为字段之一。

http-server

然后，class MultipleFigures: def __init__(self,figures): self.figures = figures def savefig(self): for fig in self.figures: fig.savefig() 函数可以返回 make fig 实例或 matplotlib.pytplot 实例。