Decorators to convert return statements to yield statements

We have this really cool tool that parses a python file into a AST and then morphs it according to some transformation and then spits out python code corresponding to the new AST. I've been playing around with it lately to come up with some transformations. One of the things I wanted to do was convert all return statements to yield expressions. That will give a good amount of coverage for yield. So this is the code I want to convert:

 def f(a,b):
    return a + b
    
def g(f, *args):
    return f(*args)

f(3,42)
g(f, 3, 42)

Now, its easy enough to change a  "return <expr>" to "yield <expr>" but now I cannot just call the function directly as f(). That will return the generator. I need to do f().next(). Now one naive approach would be to go change all the invocations of the function to f().next(). That would work except for the cases where I pass the function as a parameter to another function (like g in the example). Now when g calls calls f, it could blow up because it will get a generator instead of f's returnvalue. I could solve that also by doing replacing every f with lambda : f().next(). So the call to g become this:

 g(lambda *args: f(*args).next(), 3, 42)

That almost solves the problem except that even fs which are not references to this function (like the locally scoped f variable inside the g method) will get replaced since in a parsed AST we don't have any idea of scope (that happens during the conversion of the Python AST to the DLR AST). Now come python decorators to the rescue. In python you can wrap a function with a decorator like so:

 @dec
def f():
    pass

dec can be defined as a function. Then python will start treating f as dec(f) and any calls to f() will be semantically equivalent to dec(f)(). This is great because now python will take care of that little scoping problem for me. So now the morphed code is simply:

 def yieldwrapper(func):
    f = lambda *args, **kwargs: func(*args, **kwargs).next()
    f.__name = func.__name
    return f

@yieldwrapper
def f(a,b):
    yield a + b

@yieldwrappeer
def g(f, *args):
    yield f(*args)

f(3,42)
g(f, 3, 42)

The name of the wrapper is also set to the function's name so that it can masquerade as the real function just in case anyone's looking. Except for a few weird issues, this worked for most of the tests.