如何基于内部函数构建生成器实现Python数据传递？

需求说明 #

需要从不可修改的read_data_and_apply_func函数内部，提取每次调用process_func的参数，构建生成器传递给save_data，实现处理一个数据立即保存一个的顺序输出（如指定的期望输出所示）。仅允许修改process_func或main函数。

现有代码 #

def process_func(x): 
    # how to construct a generator out of all arguments passed 
    # to this function in the consecutive calls?
    print(f"processing {x}")
    return True


def save_func(x):
    print(f'saving {x}')



# don't edit
def read_data_and_apply_func(func):
    for x in range(5):
        func(x)


# don't edit
def save_data(generator, save_func):
    for i in generator:
        save_func(i)


def main():
    generator = ()  # generator should contain all data passed in all read_data_and_apply_func
    read_data_and_apply_func(process_func)
    save_data(generator, print)


main()

尝试过的无效代码 #

def create_generator(read_data_and_apply_func, process_func):
    def generator():
        for x in read_data_and_apply_func(generator.send):
            yield x
    return generator

def main():
    generator = create_generator(read_data_and_apply_func, process_func)
    generator_instance = generator()
    save_data(generator_instance)

期望输出 #

processing 0
saving 0
processing 1
saving 1
processing 2
saving 2
processing 3
saving 3
processing 4
saving 4

解决方案1：纯生成器协程实现严格顺序执行 #

通过修改process_func为绑定生成器的闭包，并在main中构建协调生成器，让数据处理和保存逻辑交替执行，无需额外线程即可保证严格的输出顺序。

修改后的完整代码 #

def process_func(gen):
    # 闭包绑定生成器，用于传递处理后的参数
    def inner(x):
        print(f"processing {x}")
        # 将参数x发送给生成器
        gen.send(x)
        return True
    return inner


def save_func(x):
    print(f'saving {x}')



# don't edit
def read_data_and_apply_func(func):
    for x in range(5):
        func(x)


# don't edit
def save_data(generator, save_func):
    for i in generator:
        save_func(i)


def main():
    # 定义接收参数的生成器
    def data_generator():
        while True:
            x = yield
            if x is None:
                break
            yield x
    
    # 定义协调生成器，控制处理和保存的顺序
    def coordinator():
        gen = data_generator()
        next(gen)  # 预激生成器
        # 绑定生成器到process_func
        bound_process = process_func(gen)
        # 遍历数据，每次处理后从生成器取出参数交给save_data
        for x in range(5):
            bound_process(x)
            yield next(gen)
        # 发送结束信号
        gen.send(None)
    
    # 启动保存逻辑，遍历协调生成器
    save_data(coordinator(), print)


main()

代码说明 #

1. process_func修改：改为返回闭包inner，闭包绑定了生成器gen，每次处理数据时会将参数x发送给生成器。
2. data_generator：作为数据中转的生成器，接收process_func发送的参数，并将其yield给save_data。
3. coordinator生成器：负责协调流程，预激数据生成器后，绑定process_func，然后遍历数据，每次调用处理函数后，从数据生成器取出参数并yield，让save_data执行保存操作，实现处理一个、保存一个的严格顺序。

解决方案2：队列+线程实现异步处理（顺序可能略有波动） #

如果允许轻微的顺序波动，也可以用队列+线程的方式，无需大幅修改process_func：

import queue

def process_func(x): 
    print(f"processing {x}")
    # 将参数放入队列
    q.put(x)
    return True


def save_func(x):
    print(f'saving {x}')



# don't edit
def read_data_and_apply_func(func):
    for x in range(5):
        func(x)


# don't edit
def save_data(generator, save_func):
    for i in generator:
        if i is None:
            break
        save_func(i)


def main():
    global q
    q = queue.Queue()

    # 从队列取数据的生成器
    def queue_generator():
        while True:
            item = q.get()
            yield item
    
    import threading
    # 启动保存线程
    save_thread = threading.Thread(target=save_data, args=(queue_generator(), print))
    save_thread.start()

    # 执行数据处理
    read_data_and_apply_func(process_func)

    # 发送结束信号
    q.put(None)
    save_thread.join()


main()

说明 #

• 线程调度可能导致偶尔的顺序变化，但大部分情况下会符合期望输出。
• 用None作为队列的结束信号，让save_data知道何时停止遍历生成器。

内容的提问来源于stack exchange，提问作者Johny Fa