Python函数返回生成器或普通对象的问题

Question

我定义了一个函数 f，代码如下：

def f(flag):
    n = 10
    if flag:
        for i in range(n):
            yield i
    else:
        return range(n)

但是无论 flag 是什么，f 都会返回一个生成器：

>>> f(True)
<generator object f at 0x0000000003C5EEA0>

>>> f(False)
<generator object f at 0x0000000007AC4828>

如果我对返回的对象进行遍历：

# prints normally
for i in f(True):
    print(i)

# doesn't print
for i in f(False):
    print(i)

看起来 f(False) 返回的是一个已经被遍历过的生成器。这是什么原因呢？谢谢。

Answer 1

这个生成器不会自动产生下一个值，除非你特别调用一些东西，比如

next()

我有一个叫“生成参数”的生成器，当我这样做的时候，它就正常工作了：

print(next(generate_parameters()))

Answer 2

你可以通过使用一个嵌套函数来解决这个问题，这个嵌套函数实际上会使用到 yield。

def f(flag):
    def gen():
        for i in range(n):
            yield i
    n = 10
    if flag:
        return gen()
    else:
        return range(n)

>>> f(True)
<generator object gen at 0x7f62017e3730>
>>> f(False)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

正如 Martijn 所指出的，任何包含 yield 的函数都会返回一个生成器对象。所以如果你希望在调用 f() 时，f 的主体代码能够被执行，而不是只有在遍历时才执行，你就需要采用这种方法。

标准库中的 map 方法来自 concurrent.Futures.ProcessPoolExecutor 或 concurrent.Futures.ThreadPoolExecutor，它使用这种方式来确保在调用 map 时，任务会立即被提交，而不是等到你尝试获取结果的时候。例如：

def map(self, fn, *iterables, timeout=None):
    if timeout is not None:
        end_time = timeout + time.time()

    fs = [self.submit(fn, *args) for args in zip(*iterables)]

    # Yield must be hidden in closure so that the futures are submitted
    # before the first iterator value is required.
    def result_iterator():
        try:
            for future in fs:
                if timeout is None:
                    yield future.result()
                else:
                    yield future.result(end_time - time.time())
        finally:
            for future in fs:
                future.cancel()
    return result_iterator()

Answer 3

包含 yield 语句的函数总是返回一个生成器对象。

只有当你遍历这个生成器对象时，函数里的代码才会被执行。在这之前，函数里的代码不会被执行，Python 无法知道 你只是想返回。

需要注意的是，在生成器函数中使用 return 的含义和在普通函数中是不同的; 在这种情况下，return 只是被看作是“在这里退出生成器”；返回的值会被丢弃，因为生成器只能通过 yield 表达式来产生值。

看起来你想使用 yield from 来代替：

def f(flag):
    n = 10
    if flag:
        for i in range(n):
            yield i
    else:
        yield from range(n)

yield from 需要 Python 3.3 或更高版本。

查看 yield 表达式 的文档：

在函数体中使用 yield 表达式会使该函数成为一个生成器。

当调用生成器函数时，它会返回一个称为生成器的迭代器。这个生成器会控制生成器函数的执行。执行从生成器的某个方法被调用时开始。此时，执行会进行到第一个 yield 表达式，在那里会再次暂停，将表达式列表的值返回给生成器的调用者。

遍历生成器会调用 generator.__next__() 方法，从而触发执行。

如果你想在某些情况下返回一个生成器，那么就不要在这个函数中使用 yield。你可以通过其他方式来生成生成器；例如，使用一个单独的函数，或者使用生成器表达式：

def f(flag):
    n = 10
    if flag:
        return (i for i in range(n))
    else:
        return range(n)

现在在 f 中不再使用 yield，它将不再直接产生生成器对象。相反，生成器表达式 (i for i in range(n)) 会产生它，但仅在特定条件下。