How to implement a decorated generator

Question

I have a generator:

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

I have some function that does some manipulation of the data format

def my_formatting_func(data_chunk):
    #some code
    return formated_data_chunk

What is the shortest way to create a generator that generates data_chunksin a format my_formatting_funcwithout modification my_gen?

Answer 1

Assuming Python 3.x and that the generator takes no 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.imapinstead map.