Search code examples
pythondecoratorfunctools

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

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].

Solution

  • 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.