Search code examples
pythondecoratorpython-2.xpython-decorators

How to get source code of function that is wrapped by a decorator?


I wanted to print the source code for my_func, that is wrapped by my_decorator:

import inspect
from functools import wraps

def my_decorator(some_function):
    @wraps(some_function)
    def wrapper():
        some_function()

    return wrapper

@my_decorator
def my_func():
    print "supposed to return this instead!"
    return

print inspect.getsource(my_func)

However, it returns source for wrapper instead:

@wraps(some_function)
def wrapper():
    some_function()

Is there a way for it to print the following instead?

def my_func():
    print "supposed to return this instead!"
    return

Note that the above is abstracted from a larger program. Of course we can just get rid of the decorator in this example, but that's not what I am looking for.


Solution

  • In Python 2, the @functools.wraps() decorator does not set the convenience __wrapped__ attribute that the Python 3 version adds (new in Python 3.2).

    This means you'll have to resort to extracting the original function from the closure. Exactly at what location will depend on the exact decorator implementation, but picking the first function object should be a good generalisation:

    from types import FunctionType
    
    def extract_wrapped(decorated):
        closure = (c.cell_contents for c in decorated.__closure__)
        return next((c for c in closure if isinstance(c, FunctionType)), None)
    

    Usage:

    print inspect.getsource(extract_wrapped(my_func))
    

    Demo using your sample:

    >>> print inspect.getsource(extract_wrapped(my_func))
    @my_decorator
    def my_func():
        print "supposed to return this instead!"
        return
    

    Another option is to update the functools library to add a __wrapped__ attribute for you, the same way Python 3 does:

    import functools
    
    def add_wrapped(uw):
        @functools.wraps(uw)
        def update_wrapper(wrapper, wrapped, **kwargs):
            wrapper = uw(wrapper, wrapped, **kwargs)
            wrapper.__wrapped__ = wrapped
            return wrapper
    
    functools.update_wrapper = add_wrapped(functools.update_wrapper)
    

    Run that code before importing the decorator you want to see affected (so they end up using the new version of functools.update_wrapper()). You'll have to manually unwrap still (the Python 2 inspect module doesn't go looking for the attribute); here's a simple helper function do that:

    def unwrap(func):
        while hasattr(func, '__wrapped__'):
            func = func.__wrapped__
        return func
    

    This will unwrap any level of decorator wrapping. Or use a copy of the inspect.unwrap() implementation from Python 3, which includes checking for accidental circular references.