如何在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.

Answer 3

My solution when I had to do sth like that was to create a generator on top of the generator! This is actually the idea of a decorated call. So you do,

def funca():
    while True:
        print "in funca"
        yield True

def dec(func):
    while True:
        print "in funcb"
        func.next()
        yield True

decfa = dec(funca())
decfa.next()
>>
 "in funcb"
 "in funca"

as for exactly your problem (yielding only the last value) I would do something like:

def funca():
    for i in range(1,5):
        yield i

def dec2(ff):
    try:
        while True:
            val=ff.next()
    except:
        yield val

>>>dec2(funca()).next()
4

Answer 4

Yet a simpler solution.

• Save a pointer to the initial generator as an attribute of the final one:

def dec(gen):

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

  new_gen.gen = gen
  return new_gen

• Use the pointer to the initial generator in place of the unexpectedly recursed one:

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