How to decorate a generator in python

Question

So, I defined a simple generator:

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

Basically, I want to decorate it so it returns all the values, but the last:

def dec(gen):

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

    return new_gen

Now, if I redefine gen1

@dec
def gen1(x):
    ...

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

but if I use:

some_gen = dec(gen1)

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

Why my decorator doesn't work and how can I fix it?

Answer 1

The recursive invocation of your gen1 is also subject to your decorator, so everything gets consumed by the decorator.

The simplest fix is to write the generator in non-recursive style, or to encapsulate the recursion:

Encapsulated:

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

Non-recursive:

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

Answer 2

It doesn't work due to the interaction between the decorator and recursion. Since your generator is recursive, it relies on a certain recurrence relation. By injecting a modifying decorator between the generator and the sub-generator, you are breaking that recurrence relation.

As long as @dec drops the last element, you can't make it compatible with gen1() by changing @dec alone.

You could, however, change gen1() to make it compatible with @dec:

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

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

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

The trick here is to make gen1() non-recursive, and to delegate all the work to another, undecorated, generator. The latter can be recursive.