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

Asked By: yayu

||

Source

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

Answered By: Martijn Pieters

Answer 2

Instead of yielding every potential return value, why not yield only those that actually exist? Something like

def wrap(f, arg):
    for x in f(arg):
        yield x

(actual decorator syntax, handling of positional and keyword arguments, etc. is omitted for clarity.)

Answered By: The Paramagnetic Croissant

Answer 3

For the case in comment42684128, the solution is as simple as:

(v for v in f(<args>) if filter_condition(v))

As a decorator:

def yfilter(filter_condition):
   def yfilter_p(f):
       def wrapped(*args,**kwargs):
           return (v for v in f(*args,**kwargs) if filter_condition(v))
       return wrapped
   return yfilter_p

Answered By: ivan_pozdeev

Answer 4

The existing answers don’t handle generators that yield and then return a value. For that, you need to return (yield from f()):

def dec(f):
    def g():
        return (yield from f())
    return g

@dec
def f():
   yield 'val'
   return 'done'

Answered By: Andy Jones

Answer 5

If you want to decorate using type hints, do it like this

from typing import Generator

def generate() -> Generator[int, None, None]:
    for i in range(10):
        yield i

https://stackoverflow.com/a/42531262/1812732

Answered By: John Henckel

decorating a function that yields

Question:

Answers: