How to implement decorated generator

Question

I have some generator:

def my_gen():
    while True:
        #some code
        yield data_chunk

I have some function, which makes some manipulations with data format

def my_formatting_func(data_chunk):
    #some code
    return formated_data_chunk

What the shortest way to create generator which generates data_chunks formatted by my_formatting_func without modifying my_gen?

Answer 1

Assuming Python 3.x and that the generator doesn't take any arguments (the latter is trivial to add):

def wrapper(generator):
    def _generator():
        return map(my_formatting_func, generator())
    return _generator

@wrapper
def my_gen():
    # do stuff

For 2.x, use itertools.imap instead of map.