Decorators

Lesson 30 of 31 · Part 6 · Going further · about 20 minutes

A decorator adds something to a function without changing the function's own code: a message before and after, a log of every call, a timer. You have already seen the two ideas it is built on: a function is a value you can pass around (lesson 19), and a function can be defined inside another one (lesson 18). This lesson puts them together.

Functions in, functions out

A function can take another function as an argument, and call it:

1def add_one(n):2    return n + 13 4def use_twice(func, value):5    return func(func(value))6 7print(use_twice(add_one, 5))

7

A function can also make a function and return it. The inner function remembers the variables of the outer one, even after the outer one has finished:

1def make_multiplier(factor):2    def multiply(n):3        return n * factor4    return multiply5 6double = make_multiplier(2)7triple = make_multiplier(3)8print(double(5))9print(triple(5))

10
15

1def make_multiplier(factor):2    def multiply(n):3        return n * factor4    return multiply5 6print(make_multiplier(10)(4))

Wrapping a function

Put both ideas together: a function that takes a function and returns a new one, which does something extra and calls the original in the middle. The new function is usually called wrapper.

1def announce(func):2    def wrapper():3        print("Wick picks up the ladder")4        func()5        print("Wick puts the ladder away")6    return wrapper7 8def light_lamp():9    print("The lamp is lit")10 11light_lamp = announce(light_lamp)12light_lamp()

Wick picks up the ladder
The lamp is lit
Wick puts the ladder away

announce is a decorator. The line light_lamp = announce(light_lamp) replaces the name light_lamp with the wrapper, and the wrapper still holds on to the original function as func. Press Step to follow the call from wrapper into the original and back.

The @ syntax

Writing a function's name three times, as in light_lamp = announce(light_lamp), is clumsy. Python has a shortcut: put @announce on the line above the def. It means exactly the same thing.

1def announce(func):2    def wrapper():3        print("Wick picks up the ladder")4        func()5        print("Wick puts the ladder away")6    return wrapper7 8@announce9def light_lamp():10    print("The lamp is lit")11 12light_lamp()

Wick picks up the ladder
The lamp is lit
Wick puts the ladder away

Now anyone reading light_lamp sees straight away that it is decorated. The decorator runs once, when the def is reached, not every time the function is called:

1def announce(func):2    print("Wrapping", func.__name__)3    return func4 5@announce6def light_lamp():7    print("Lit")8 9print("Ready")

Functions with arguments

The wrapper above takes no arguments, so it only works for functions that take none. Decorate add(a, b) with it and calling add(2, 3) fails with a TypeError: the wrapper takes 0 arguments but got 2. The fix is *args and **kwargs from lesson 17: the wrapper accepts any arguments and passes them all on. It also returns what the original returned.

1def log_call(func):2    def wrapper(*args, **kwargs):3        print("Calling", func.__name__, args, kwargs)4        result = func(*args, **kwargs)5        print("Returned", result)6        return result7    return wrapper8 9@log_call10def add(a, b):11    return a + b12 13@log_call14def greet(name, greeting="Hello"):15    return f"{greeting}, {name}!"16 17add(2, 3)18greet("Wick", greeting="Good evening")

Calling add (2, 3) {}
Returned 5
Calling greet ('Wick',) {'greeting': 'Good evening'}
Returned Good evening, Wick!

This is the shape almost every decorator has. One decorator now works for any function, whatever arguments it takes. Leave out return result and every decorated function would return None.

A timing decorator

A classic use: measure how long a function takes, without touching its code. time.perf_counter() (lesson 25) is read before and after the call, and the difference is the time in seconds.

1import time2 3timings = {}4 5def timed(func):6    def wrapper(*args, **kwargs):7        start = time.perf_counter()8        result = func(*args, **kwargs)9        elapsed = time.perf_counter() - start10        timings[func.__name__] = elapsed11        return result12    return wrapper13 14@timed15def total_ore(carts):16    return sum(range(carts))17 18print(total_ore(10_000))19print(list(timings))20print(timings["total_ore"] >= 0)

49995000
['total_ore']
True

In your own programs the wrapper would simply print the time, for example with f"{elapsed:.4f} s". Here it stores the time in a dictionary instead, because the number is different on every run and every computer. What you can see: the result is passed through unchanged, and the time was recorded under the function's name.

functools.wraps

A wrapper has a small side effect: the decorated function now reports the wrapper's name and docstring, not its own.

1def log_call(func):2    def wrapper(*args, **kwargs):3        return func(*args, **kwargs)4    return wrapper5 6@log_call7def add(a, b):8    """Add two numbers."""9    return a + b10 11print(add.__name__)12print(add.__doc__)

wrapper
None

That confuses error messages, logs and anyone using help(). The functools module has a decorator for exactly this problem: put @functools.wraps(func) on the wrapper and it copies over the original's name and docstring.

1import functools2 3def log_call(func):4    @functools.wraps(func)5    def wrapper(*args, **kwargs):6        return func(*args, **kwargs)7    return wrapper8 9@log_call10def add(a, b):11    """Add two numbers."""12    return a + b13 14print(add.__name__)15print(add.__doc__)

add
Add two numbers.

functools.wraps(func) is a call that returns a decorator, which is why it has parentheses. Make it a habit: every wrapper you write gets one.

Decorators that come with Python

You already met @dataclass in lesson 28. Python has a few more decorators for methods in classes:

1class Lamp:2    def __init__(self, oil):3        self.oil = oil4 5    @property6    def is_lit(self):7        return self.oil > 08 9    @staticmethod10    def max_oil():11        return 1012 13    @classmethod14    def full(cls):15        return cls(cls.max_oil())16 17lamp = Lamp.full()18print(lamp.oil)19print(lamp.is_lit)20print(Lamp.max_oil())

10
True
10

Lamp.full() is called on the class, not on a lamp, and builds a lamp with a full tank. lamp.is_lit has no parentheses, yet it runs the method every time it is read.

Practice

Write the code, then press Check. A hint costs nothing; the answer is there when you need it.

Ding!
Ding!

1def twice(func):2    pass3 4@twice5def ring():6    print("Ding!")7 8ring()

GOOD EVENING, WICK

1def shout(func):2    pass3 4@shout5def greet(name, greeting="Hello"):6    return f"{greeting}, {name}"7 8print(greet("Wick", greeting="Good evening"))

light_lamp

1import functools2 3def announce(func):4    def wrapper():5        print("Here we go")6        func()7    return wrapper8 9@announce10def light_lamp():11    print("Lit")12 13print(light_lamp.__name__)

What you learned

• Functions can take functions as arguments and return new functions; an inner function remembers the outer one's variables.
• A decorator takes a function and returns a wrapper that adds something before or after calling it.
• @decorator above a def is short for name = decorator(name), and runs once, when the function is defined.
• Give the wrapper *args, **kwargs, pass them on and return the result, so it works for any function.
• @functools.wraps(func) keeps the original name and docstring; @property, @staticmethod, @classmethod and @dataclass (from dataclasses) come with Python.