python - 如何在python中装饰一个生成器

Question

所以，我定义了一个简单的生成器:

def gen1(x):
    if x <= 10:
        yield x
        for v in gen1(x + 1):
            yield v

基本上，我想装饰它以返回所有值，但最后一个值:

def dec(gen):

    def new_gen(x):
        g = gen(x)
        value = g.next()
        for v in g:
            yield value
            value = v

    return new_gen

现在，如果我重新定义 gen1

@dec
def gen1(x):
    ...

for i in gen1(1):
    print i    # Nothing printed

但如果我使用:

some_gen = dec(gen1)

for i in some_gen(1):
    print i    # Prints 1 to 9, as needed

为什么我的装饰器不起作用，我该如何解决？

Answer 1

您的 gen1 的递归调用也受制于您的装饰器，因此装饰器会消耗所有内容。

最简单的解决方法是以非递归方式编写生成器，或者封装递归:

封装:

@dec
def gen1(x):
    def inner(x):
        if x <= 10:
            yield x
            for v in inner(x + 1):
                yield v
    return inner(x)

非递归:

@dec
def gen1(x):
    for v in range(x, 11):
        yield v