如何包装或嵌入生成器？

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. 错误表明“文件”仍然是一个生成器，无论输入是单个目录还是它的列表。为什么它仍然是发电机？
2. 是否可以在函数中包装或嵌入生成器？如果是这样，如何使这项工作？

Answer 1

你每次都让helper()：

yield helper(dir)

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

在 Python 3.3 及更高版本中，请改用 yield from：

yield from helper(dir)

这个委派控制权给另一个生成器。来自Yield expressions 文档：

当使用yield from <expr> 时，它将提供的表达式视为子迭代器。该子迭代器生成的所有值都直接传递给当前生成器方法的调用者。

在较旧的 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'):
    # ...