Edgar Navasardyan
Edgar Navasardyan

Reputation: 4501

Python - decorator manipulating attributes of target functions

My goal is to come up with a decorator that would accept functions with arbitrary number of params and change the target function's attributes. One example, I need an attribute in the target function which would contain execution length in seconds. I have tried to utilize functool's wraps method to preserve the attributes, but with no success. It only seems to preserve the initial attributes like __doc__ and __name__, but I am not able to assign new value to non-existant attribute. I know that I could assign the duration attribute to wrapper before returning it, intead of func. But this is suitable only for those cases when we do not have to count for what happens inside the func, like in my case. So, here's my code:

import time
from functools import wraps

def my_dec(param):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            start = time.time()
            rs = func(*args, **kwargs)
            func.duration = time.time() - start
            return rs
        wrapper.a = 'a'
        return wrapper
    return decorator


@my_dec('Dec param')
def target_func(x):
    return x.upper()
target_func('test value')
print(target_func.duration)

Upvotes: 3

Views: 244

Answers (2)

Jiří Baum
Jiří Baum

Reputation: 6930

Alternative answer, if you want to avoid referencing the wrapper function within its own definition:

import time
from functools import wraps

def my_dec(param):
    def decorator(func):
        durations = []
        @wraps(func)
        def wrapper(*args, **kwargs):
            start = time.time()
            rs = func(*args, **kwargs)
            durations.append(time.time() - start)
            return rs
        wrapper.durations = durations
        return wrapper
    return decorator


@my_dec('Dec param')
def target_func(x):
    return x.upper()
target_func('test value')
print(target_func.durations)

Upvotes: -1

Jiří Baum
Jiří Baum

Reputation: 6930

Assigning to wrapper.duration works...

import time
from functools import wraps

def my_dec(param):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            start = time.time()
            rs = func(*args, **kwargs)
            wrapper.duration = time.time() - start
            return rs
        return wrapper
    return decorator


@my_dec('Dec param')
def target_func(x):
    return x.upper()
target_func('test value')
print(target_func.duration)

Upvotes: 3

Related Questions