Understanding Python Function Decorators

Posted on In ,

Decorators let you wrap extra behavior around a function without touching its source.

Why decorators?

Suppose we start with:

1
2
3
4
import random

def func_a():
    print("I'm func_a, get a random number %s" % random.random())

We want to print an extra line whenever func_a runs. Editing the function directly works, but it does not scale when the same tweak applies to many functions—or if we cannot edit the source.

First attempts

Wrapping the function manually:

1
2
3
def new_func_a():
    print("I'm new code")
    func_a()

Works but requires one wrapper per function. Passing the original function into a helper is cleaner:

1
2
3
4
5
def add_new_code(func):
    print("I'm new code")
    return func

new_func_a = add_new_code(func_a)

The helper runs immediately, so the extra line prints during assignment, not when the function executes. We need to return a callable that delays the work.

Proper wrapper

1
2
3
4
5
6
7
8
9
10
11
12
from functools import wraps

def add_new_code(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print("I'm new code")
        return func(*args, **kwargs)
    return wrapper

@add_new_code
def func_a():
    print("I'm func_a, get a random number %s" % random.random())

@add_new_code is syntactic sugar for func_a = add_new_code(func_a). functools.wraps preserves metadata such as __name__.

Passing arguments to decorators

Decorators that accept parameters return another decorator:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from functools import wraps

def add_new_code(msg):
    def decorate(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            print("I'm new code")
            print("Get a arg: %s" % msg)
            return func(*args, **kwargs)
        return wrapper
    return decorate

@add_new_code('wrapper for func_a')
def func_a():
    print("I'm func_a, get a random number %s" % random.random())

Calling add_new_code('wrapper for func_a') returns decorate, which then wraps func_a.

Class-based decorators

Any callable can act as a decorator. Implement __call__ on a class to wrap functions and maintain state:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import random
from functools import wraps

class Profiled:
    def __init__(self, func):
        self.func = func
        self.call_count = 0
        wraps(func)(self)

    def __call__(self, *args, **kwargs):
        self.call_count += 1
        print("I'm new code")
        print("new func running count: %s" % self.call_count)
        return self.func(*args, **kwargs)

@Profiled
def func_a():
    print("I'm func_a, get a random number %s" % random.random())

Examples

timethis measures execution time:

1
2
3
4
5
6
7
8
9
10
11
12
import time
from functools import wraps

def timethis(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(func.__name__, end - start)
        return result
    return wrapper

logged adds logging:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from functools import wraps
import logging

def logged(level, name=None, message=None):
    def decorate(func):
        logname = name if name else func.__module__
        log = logging.getLogger(logname)
        logmsg = message if message else func.__name__

        @wraps(func)
        def wrapper(*args, **kwargs):
            log.log(level, logmsg)
            return func(*args, **kwargs)
        return wrapper
    return decorate