How to extend the functionality of the functools.wraps decorator?

Question

I'd like to create a new decorator to use in place of @wraps(f) that does whatever magic @wraps(f) would do as well as something else. How would I do that?

Specifically, I've got several decorators that are of the form:

def decorator(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        # does some stuff in here
        return f(*args, **kwargs)
    if not hasattr(wrapper, '_first_line'):
        wrapper._first_line = inspect.getsourcelines(f)[1]
    return wrapper

It seems like I should be able to create a decorator like @wraps_with_first_line(f) which will do all that @wraps(f) is doing as well as if not hasattr(wrapper, '_first_line'): wrapper._first_line = inspect.getsourcelines(f)[1].

Answer 1

If what you want to add isn't already an attribute of the wrapped object, then you can use this:

def wraps_with_first_line(f):
    def wrap(wrapper):
        wrapper = wraps(f)(wrapper)
        if not hasattr(wrapper, '_first_line'):
            wrapper._first_line = inspect.getsourcelines(f)[1] 
        return wrapper
    return wrap

If it is already an attribute of the wrapped object, use Sven's method.

Answer 2

You should rather follow the good practice of adding a __wrapped__ attribute pointing to the wrapped function than adding single attributes of that wrapped function. New versions of functools.wraps() do this automatically, but if you are using an older version of Python than 3.2, you can als easily extend wraps() to add __wrapped__:

def my_wraps(wrapped, **kwargs):
    def decorator(wrapper):
        functools.update_wrapper(wrapper, wrapped, **kwargs)
        wrapper.__wrapped__ = wrapped
    return decorator

Edit: Here's a function the extracts the original function from a possibly multiply decorated function:

def orig_function(f):
    try:
        while True:
            f = f.__wrapped__
    except AttributeError:
        return f