如何包装或嵌入生成器？

Question

我试图提供一个统一的界面，用于从单个目录或目录列表中检索所有文件。

def get_files(dir_or_dirs):
    def helper(indir):
        file_list = glob.glob("*.txt")
        for file in file_list:
            yield file

    if type(dir_or_dirs) is list:
        # a list of source dirs
        for dir in dir_or_dirs:
            yield helper(dir)
    else:
        # a single source dir
        yield helper(dir_or_dirs)

def print_all_files(file_iter):
    for file in file_iter:
        print(file)        # error here!

问题：

错误说 'file' 仍然是一个生成器，无论输入是单个目录还是目录列表。为什么还是发电机？
是否可以在函数中包装或嵌入生成器？如果是这样，如何使这项工作？

Answer 1

您每次都获得 helper()：

yield helper(dir)

但是helper()本身就是一个发电机。

在 Python 3.3 和更新版本中，使用 yield from 代替：

yield from helper(dir)

这委托控制另一台发电机。来自 Yield expressions 文档：

When yield from <expr> is used, it treats the supplied expression as a subiterator. All values produced by that subiterator are passed directly to the caller of the current generator’s methods.

在旧的 Python 版本中，包括 Python 2.x，使用另一个循环：

for file in helper(dir):
    yield file

有关 yield from 功能的更多信息，请参阅 PEP 380 -- Syntax for Delegating to a Subgenerator。

并不是说您真的需要辅助函数，它只是循环遍历 glob.glob() 结果，您可以直接执行此操作。

您还需要更正您的功能才能实际使用 indir；目前您忽略了该参数，因此您只能从 当前工作目录 .

获取文本文件

接下来，您想使用 glob.iglob() 而不是 glob.glob() 来对 os.scandir() 进行惰性求值，而不是一次将所有结果加载到内存中。我只是将一个非列表 dir_or_dirs 值变成一个列表，然后只使用一个循环：

import glob
import os.path

def get_files(dirs):
    if not isinstance(dirs, list):
        # make it a list with one element
        dirs = [dirs]

    for dir in dirs:
        pattern = os.path.join(dir, '*.txt')
        yield from glob.iglob(pattern)

现在，我不再使用字符串或列表形式的单个参数，而是使用可变数量的参数，使用 *args 参数语法：

def get_files(*dirs):
    for dir in dirs:
        pattern = os.path.join(dir, '*.txt')
        yield from glob.iglob(pattern)

这可以用 0 个或多个目录调用：

for file in get_files('/path/to/foo', '/path/to/bar'):
    # ...

如何包装或嵌入生成器？

How to wrap or embed generators?

python

yield

generator