python 函数作为可选的生成器

Question

假设我有一个函数，我想选择是否返回结果。这很容易编码：

def foo(N, is_return=False):
    l = []
    for i in range(N):
        print(i)
        if is_return:
            l.append(i)
    if is_return:
        return l

但是现在可以说我希望该函数成为一个生成器。我会写这样的东西：

def foo_gen(N, is_return=False):
    for i in range(N):
        print(i)
        if is_return:
            yield i

所以大概 when is_returnis Falsethenfoo_gen只是一个没有返回值的函数， when is_returnisTrue foo_gen是一个生成器，我希望有两个不同的调用：

In [1]: list(foo_gen(3, is_return=True))
0
1
2
Out[2]: [0, 1, 2]

因为当它是一个生成器并且您必须遍历产生的值时，并且：

>>> In [2]: foo_gen(3)
0
1
2

因为当它不是生成器并且它只是具有副作用并且您不必遍历它时。但是，后一种行为不起作用，而只是返回生成器。你不能从中得到任何东西：

In [3]: list(foo_gen(3, is_return=False))
0
1
2
Out[3]: []

但这并不是很好，并且对于那些不希望必须遍历任何内容以产生副作用的 API 用户来说是令人困惑的。

反正有没有In [2]在函数中做出行为？

Answer 1

为此，您需要包装foo_gen在另一个函数中，该函数返回生成器或对其本身进行迭代，如下所示：

def maybe_gen(N, is_return=False):
    real_gen = foo_gen(N)
    if is_return:
        for item in real_gen:
            pass
    else:
        return real_gen

def foo_gen(N):
    for i in range(N):
        print(i)
        yield i

>>> list(maybe_gen(3))
0
1
2
[0, 1, 2]
>>> maybe_gen(3, is_return=True)
0
1
2
>>> 

原因是yield函数中任何地方的出现使它成为生成器函数。没有办法让函数在调用时决定它是否是生成器函数。相反，您必须有一个非生成器函数，它在运行时决定是返回生成器还是其他东西。

也就是说，这样做很可能不是一个好主意。你可以看到maybe_genwhen is_returnis True做什么是完全微不足道的。它只是迭代生成器而不做任何事情。这尤其愚蠢，因为在这种情况下，生成器本身除了打印之外什么都不做。

最好让函数 API 保持一致：要么总是返回一个生成器，要么永远不做。一个更好的主意是只有两个函数foo_gen，即生成器，print_gen或者无条件打印它的东西。如果你想要生成器，你可以调用foo_gen. 如果您只是想打印它，则print_gen改为调用，而不是将“标志”参数传递给foo_gen.

关于你最后的评论：

但这并不是很好，并且对于那些不希望必须遍历任何内容以产生副作用的 API 用户来说是令人困惑的。

如果 API 指定函数返回一个生成器，用户应该期望必须对其进行迭代。如果 API 说它不返回一个生成器，用户不应该期望必须对其进行迭代。API 应该只说一个或另一个，这将使用户清楚地知道期望什么。更令人困惑的是有一个笨拙的 API，它告诉用户他们必须传递一个标志来确定他们是否获得了生成器，因为这使用户的期望变得复杂。