如何通过装饰器将Python生成器改造为可复用可迭代对象

普通生成器函数调用后会返回一次性迭代器，遍历完成后就会耗尽，无法重复迭代。要实现类似内置range的可复用效果，只需要编写一个通用装饰器，让原函数返回自定义可迭代对象，每次迭代时重新生成全新的生成器实例即可。

装饰器实现 #

import functools

def reusable(func):
    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        class IterableFactory:
            def __iter__(self):
                # 每次迭代都重新执行原生成器函数，返回全新迭代器
                return func(*args, **kwargs)
        return IterableFactory()
    return wrapped

使用方式 #

只需要在定义的生成器函数上加上这个装饰器即可：

@reusable
def myrange(n):
    for i in range(n):
        yield i

效果验证 #

执行测试代码，可以看到两次调用list()都能得到完整结果：

x = myrange(3)
print(list(x)) # [0, 1, 2]
print(list(x)) # [0, 1, 2]

实现原理 #

• 未加装饰器时，调用myrange(3)会直接返回生成器迭代器，内部维护单次迭代的游标位置，遍历到末尾后就会持续抛出StopIteration，后续迭代无法拿到值
• 加装饰器后，调用myrange(3)返回的是IterableFactory类的实例，这个实例本身是可迭代对象（实现了__iter__方法），但不是迭代器，不会保存单次迭代的状态
• 每次对该实例做迭代操作（比如传入list()、用于for循环）时，__iter__方法都会用最初传入的参数重新调用原生成器函数，得到一个从头开始的全新生成器，因此可以反复遍历得到完整结果
• 装饰器中使用functools.wraps是为了保留原生成器函数的元信息（如函数名、文档字符串），避免原函数属性被覆盖。

内容的提问来源于stack exchange，提问作者Ilya V. Schurov