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

Question

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.

Answer 1

I had a need to retrieve the original function that has been decorated multiple times. I wrote the following code that will try to unwrap until it gets to a function that hasn't been decorated.

def extract_wrapped(decorated):
    closure = decorated.__closure__
    if closure:
        for cell in closure:
            if isinstance(cell.cell_contents, FunctionType) and cell.cell_contents.__closure__ is None:
                return cell.cell_contents
            elif isinstance(cell.cell_contents, FunctionType):
                return extract_wrapped(cell.cell_contents)
    return decorated

You can use the extract_wrapped function to get the original function that has been decorated. For example:

import inspect
from types import FunctionType


def decorator(func):
    def wrapper():
        print("before")
        func()
        print("after")
    return wrapper


@decorator
@decorator
def my_function():
    print("func")


print("Decorated:")
my_function()
# before
# before
# func
# after
# after


def extract_wrapped(decorated):
    closure = decorated.__closure__
    if closure:
        for cell in closure:
            if isinstance(cell.cell_contents, FunctionType) and cell.cell_contents.__closure__ is None:
                return cell.cell_contents
            elif isinstance(cell.cell_contents, FunctionType):
                return extract_wrapped(cell.cell_contents)
    return decorated


original = extract_wrapped(my_function)

print("\nOriginal source:")
print(inspect.getsource(original))
# @decorator
# @decorator
# def my_function():
#     print("func")

print("\nOriginal:")
original()
# func

Answer 2

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.

Answer 3

As Martijn Pieters points out in his answer, the Python 2 @functool.wraps() decorator doesn't define a __wrapped__ attribute, which would make doing what you want to do very easy. According to the documentation I read, even though it was added in Python 3.2, there was a bug in the ways it was sometimes handled until version 3.4 was released—so the code below uses v3.4 as the cut-off for defining a custom wraps() decorator.

Since from its name it sounds like you have control over my_decorator(), you can workaround the issue by defining you're own wraps-like function, rather than extracting the original function from the closure, as shown in his answer. Here's how to do it (which works in Python 2 and 3):

^{(As Martijn also points out, you could monkey-patch the change in by overwriting the functools.wraps module attribute, which would make the change also affect other modules that use functools instead of only the one where it's defined.)}

import functools
import inspect
import sys

if sys.version_info[0:2] >= (3, 4):  # Python v3.4+?
    wraps = functools.wraps  # built-in has __wrapped__ attribute
else:
    def wraps(wrapped, assigned=functools.WRAPPER_ASSIGNMENTS,
              updated=functools.WRAPPER_UPDATES):
        def wrapper(f):
            f = functools.wraps(wrapped, assigned, updated)(f)
            f.__wrapped__ = wrapped  # set attribute missing in earlier versions
            return f
        return wrapper

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.__wrapped__))

Output:

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