0

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

我有一个make_fig过去只生成一个图形的函数，我会像fig在许多其他函数中一样调用它，然后将其保存为fig.savefig.

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

fig = make_fig

for f in fig:
   f.savefig

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

我可以以某种方式修改该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')
4

2 回答 2

1

只需实现可以处理这两种情况的自己的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.jpgtest2_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')
于 2021-01-17T21:41:19.280 回答
1

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

class MultipleFigures:

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

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

然后该make fig函数可以返回一个matplotlib.pytplot实例或一个MultipleFigures实例。

于 2021-01-17T22:09:51.893 回答