Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Decorating a generator in Python: call some method in between yields

I found some very useful information about decorating generator functions in Python here using yield from. For example:

def mydec(func):
    def wrapper(*args, **kwargs):
        print(f'Getting values from "{func.__name__}"...')
        x = yield from func(*args, **kwargs)
        print(f'Got value {x}')
        return x
    return wrapper

@mydec
def mygen(n):
    for i in range(n):
        yield i

However, this seems to only allow for adding decorated behaviors at the beginning and end of the generator's lifetime:

>>> foo = mygen(3)
>>> x = next(foo)
Getting values from "mygen"...
>>> x
0
>>> x = next(foo)
>>> x
1
>>> x = next(foo)
>>> x
2
>>> x = next(foo)
Got value None
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>> x
2

However I am interested in using the decorator to implement some behavior every time the generator yields. However, the decorator should not modify the values that are gotten from the generator. That is, for example, I'd like to have the output:

>>> foo = mygen(3)
>>> x = next(foo)
Getting values from "mygen"...
Got value 0
>>> x
0
>>> x = next(foo)
Got value 1
>>> x
1
>>> x = next(foo)
Got value 2
>>> x
2

So, a call to print occurs with each yield, however the yielded values remain unchanged.

Is this possible?

like image 726
corvus Avatar asked Oct 30 '25 22:10

corvus


1 Answers

yield from is for coroutine stuff. You're not doing coroutine stuff. Just iterating the generator:

def mydec(func):
    def wrapper(*args, **kwargs):
        print(f'Getting values from "{func.__name__}"...')
        gen = func(*args, **kwargs)
        for value in gen:
            print(f'got value {value}')
            yield value
    return wrapper
like image 127
wim Avatar answered Nov 01 '25 13:11

wim