python - 如何包装或嵌入生成器？

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!

问题:

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

Answer 1

您正在屈服 helper()每次:

yield helper(dir)

但是helper()本身就是一个生成器。

在 Python 3.3 和更新版本中，使用 yield from相反:

yield from helper(dir)

这委托(delegate)控制给另一个生成器。来自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'):
    # ...