decorating a function that yields

Question

Is it possible, and if so, advisable, and if so, what would be the recommended method for decorating a function that yields a value?

For example, consider this imaginary example I made up

def foobar_creator(func):
    def wrapped(**kwargs):
       res = func(**kwargs)
       flag = True
       for k,v in kwargs:
           if res % v == 0:
               flag = False
               yield k
       if flag: 
            yield res
     return wrapped

@foobar_creator
def generic_yielder(**kwargs):
     for i in xrange(sys.maxint):
           yield i

for i in generic_yielder(foo=3, bar=5, foobar=15):
      print i

Answer 1

A generator function, when called, returns an iterator object. If your decorator is itself a generator too, you'll need to loop over the wrapped result:

def foobar_creator(func):
    def wrapped(**kwargs):
        gen = func(**kwargs)
        flag = True
        for k, v in kwargs:
            if res % v == 0:
                flag = False
                yield k
        if flag:
            for res in gen:
                yield res
    return wrapped

If you are using Python 3.3 or up, you can use delegation to hand control the wrapped generator, by using yield from:

if flag:
    yield from gen