Command Palette

Search for a command to run...

Hectal
PHASE 9Intermediate ~26 min· topic 7 of 9

Topic 9.7

Decorators with Arguments and Class Decorators

In one line

A decorator that takes its own settings, such as @retry(attempts=3), needs three levels of functions: one to take the settings, one to take the function, one to wrap the call. Decorators can also be written as classes, applied to methods, and applied to whole classes.

Think of it like this

Ordering a coffee. @decorator is like saying "a coffee" and getting the standard one. @decorator(sugar=2) is like saying "a coffee with two sugars": you first tell the barista your settings, and the barista then makes a customised coffee. In code, a decorator with arguments is a function that returns a decorator.

Words you'll meet

New words in this topic, in plain English. Come back here whenever one feels fuzzy.

Decorator factory
A function that takes settings and returns a decorator. repeat(3) is a decorator factory.
Three-level nesting
The shape of a decorator with arguments: settings function, decorator function, wrapper function.
Keyword-only parameter
A parameter after * that must be passed by name, like label="API".
Callable
Anything you can call with brackets: a function, a method, or an object with __call__.
__call__
The special method that lets an object be called like a function.
Class decorator
A decorator applied to a class, receiving the class and returning a class.
Registry
A dictionary or list in which decorated functions or classes record themselves.

Step by step

01The problem: a decorator that needs a setting

A plain decorator receives exactly one thing: the function. If you want times=3, there is nowhere to put it. The answer is one more function on the outside whose only job is to catch the setting and return the real decorator.

02Three levels

Read it from the outside in. repeat(times) is called with the setting and returns decorator. decorator(func) is called with the function and returns wrapper. wrapper(*args, **kwargs) runs on every call and can use times and func from the enclosing scopes.

main.pywhole filepython
import functools

def repeat(times):                      # level 1: the settings
    def decorator(func):                # level 2: the function
        @functools.wraps(func)
        def wrapper(*args, **kwargs):   # level 3: every call
            result = None
            for _ in range(times):
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator

@repeat(times=3)
def hello(name):
    print("hello", name)
Three levelsdiagram
Rendering diagram…

03A configurable retry

The retry from Topic 9.6 becomes much more useful with settings: how many attempts, and which exceptions count as temporary. The default exceptions=(ConnectionError,) is a tuple, so except exceptions can catch several types at once.

The example calls a fake function that fails twice with ConnectionError, then succeeds. No real waiting happens, so the output is fixed.

04Optional brackets

To allow both @tag and @tag(label="X"), check whether the first argument is a function. If it is, you were called the short way and decorate right away. If not, you were called with settings, so return the decorator for the function that will arrive next.

main.pywhole filepython
def tag(_func=None, *, label="LOG"):
    def decorator(func):
        ...   # build and return the wrapper
    if _func is not None:        # used as @tag
        return decorator(_func)
    return decorator             # used as @tag(label="X")

05Decorating a method

Inside a class, @log_method above def reply(self, text) wraps the plain function before it becomes a method. Python then binds self when you call bot.reply(...), so the wrapper receives self as its first argument. Naming it self in the wrapper lets the decorator use it (type(self).__name__).

06A decorator as a class

Counted(work) creates an instance that stores the function and a counter. Calling work() calls the instance's __call__, which increments the counter and calls the function. The state lives in self.calls, with no nonlocal needed (nonlocal lets an inner function change an outer variable, Topic 5.5).

functools.update_wrapper(self, func) is the class version of wraps: it copies the name and docstring onto the instance.

07A class decorator

@register receives the class itself, adds registered = True to it and returns it. The class is the same class as before, with one more attribute. Registries work the same way: REGISTRY[cls.__name__] = cls; return cls.

08When not to use a decorator

If a decorator needs three settings and nested levels, a plain function that takes a function (retry_call(func, attempts=3)) or a small class may read better. Decorators are best when the same wrapping is applied to many functions and the behaviour is easy to describe in a name.

Try it yourself

  1. 1

    Count the levels

    In the first example, add a print("level 1") in repeat, a print("level 2") in decorator and a print("level 3") in wrapper. Run it. Levels 1 and 2 print once, when the def is executed. Level 3 prints on every call.

  2. 2

    Retry only some errors

    Change the fake function to raise ValueError("bad input") instead. Because retry only catches ConnectionError, the error is raised at once, with no retry. That is the intended behaviour for bugs.

Code & diagrams

Decorators with arguments New tab

Change times=3 to times=2, and attempts=4 to attempts=2 to see the retry give up.

Sign in to run this example in your browser.

Expected output

hello Ada
hello Ada
hello Ada
attempt 1/4 failed: timeout
attempt 2/4 failed: timeout
ok
Optional brackets, class decorators and methods New tab

Add a second decorated method to Bot and call it. Then apply @tag(label='DB') to a new function.

Sign in to run this example in your browser.

Expected output

[LOG] a
[API] b
work called 2 times, name: work
Plugin.registered: True
Bot.reply called
cba

Break it on purpose

Errors are the best teachers. Make each change, read the error, guess what went wrong, then reveal the answer.

Break #1

Use a decorator factory without brackets

Write @repeat instead of @repeat(3).

main.pywhole filepython
def repeat(times):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for _ in range(times):
                func(*args, **kwargs)
        return wrapper
    return decorator

@repeat
def hello():
    print("hello")

hello()
terminal
$ python main.py
── what you'll see ──
Traceback (most recent call last):
File "main.py", line 13, in <module>
hello()
~~~~~^^
TypeError: repeat.<locals>.decorator() missing 1 required positional argument: 'func'

Myth vs fact

Myth

A decorator with arguments is just a decorator with extra parameters.

Fact

It is a function that returns a decorator. That is why it has one more level of nesting.

Myth

Class decorators and decorators written as classes are the same thing.

Fact

A class decorator is a decorator applied *to a class*. A decorator written as a class is a class used *as* a decorator (usually on functions). They are different ideas.

Myth

A class-based decorator works on methods just like a function-based one.

Fact

An instance of a plain class does not bind self when used on a method. The function-based form (or a __get__ method) is needed.

Pro corner

Extra depth for experienced readers. New to this? Skip it for now and come back later.

  • ▸

    Decorator arguments are evaluated at definition time, in the order @a(x), @b(y) from top to bottom, and the resulting decorators are *applied* bottom to top. Side effects in the factory (like registering the function in a table) therefore happen once, at import.

  • ▸

    For a decorator factory with a mutable default (for example exceptions=[]), the default is shared between all uses, the same trap as mutable default arguments in functions (Topic 5.2). Use tuples or None.

  • ▸

    Registries built with decorators are the standard plug-in mechanism: web frameworks use @app.get("/path"), test runners use @pytest.fixture, and agent frameworks use a decorator to turn a Python function into a tool the model can call, reading its name, docstring and type hints. That is why wraps and annotations matter.

  • ▸

    A class decorator that returns a different object (a wrapper class or a function) breaks isinstance checks and subclassing. Return the same class, modified, unless you have a good reason. Even the standard library has an exception: @dataclass(slots=True) has to build a new class (a class cannot gain __slots__ after creation), so any reference to the original class taken inside the decorated class body is to a different object from the final class. For hooks that run when a class is subclassed, __init_subclass__ is the cleaner tool.

Remember this

  1. 1

    @repeat(times=3) is a *call* to repeat(3) that must return a real decorator, and that decorator is then applied to the function. So you need three levels: the outer function takes the settings (times), the middle function takes the function being decorated, and the inner function is the wrapper that runs on every call. Python evaluates repeat(times=3) first, then applies the result to the function.

  2. 2

    The most common mistake is to forget the brackets or the extra level. @repeat without brackets passes the function where times should be, and the failure appears later, when you call the function, as an odd error about a missing func argument. If a decorator has settings, always write @decorator(...), even with no arguments, unless you make the brackets optional.

  3. 3

    Optional arguments. Some decorators work both as @tag and as @tag(label="API"). The standard trick: make the decorator's first parameter the function with a default of None, and the other parameters keyword-only (after a bare *, Topic 5.4). If the function arrives, decorate it right away; if not, return the decorator for the function that comes next.

  4. 4

    Decorating methods. A method's first parameter is self, but a wrapper using *args, **kwargs forwards it unchanged, so the same decorator works on methods. If the wrapper needs self (for example to log the class name), name it as the first parameter. Decorators like @staticmethod, @classmethod and @property (Topics 8.3 and 8.4) are decorators too; put your own decorator below them (closer to the def) so they receive a plain function.

  5. 5

    Decorators written as classes. Any callable can be a decorator, including a class whose __init__ takes the function and whose __call__ runs the wrapped call. This is handy when the decorator needs state, such as a counter or a cache, because the state lives in attributes. Use functools.update_wrapper(self, func) in __init__ so the instance copies the function's name and docstring, just as wraps does. One catch: a class instance used on a method does not receive self unless it also implements __get__, so function-based decorators are safer for methods.

  6. 6

    Class decorators. A decorator can take a class and return it (usually after changing it): @register can add the class to a registry, or add an attribute. @dataclass (Topic 8.8) is exactly this: it reads the class and adds methods such as __init__ and __repr__. Many frameworks use class decorators for registration and configuration. For anything more complex than adding or checking attributes, a base class or __init_subclass__ may be clearer.

Explain it without notes

01

How does a decorator that takes arguments work? Describe the three levels.

02

How would you write a decorator that works both with and without brackets?

03

What is the difference between a class decorator and a decorator written as a class?

04

Why can a class-based decorator be awkward on methods?

05

In what order are the arguments of stacked decorator factories evaluated, and when does each part of a three-level decorator run?

Practice

01

Write a decorator factory prefix(text) that adds text in front of the string a function returns. Use @prefix(">> ") on a function returning "hi" and print the result.

02

Write a class decorator add_hello that adds a method hello to a class, returning "hello from " + the class name. Apply it to class Dog and call Dog().hello().

03

Write a decorator factory limit_calls(n) that lets a function run only n times and raises RuntimeError("limit reached") on the next call. Show it with n=2.

Trade-offs

  • ↔

    Power versus readability. Three levels of nesting are hard to read. Keep the inner wrapper small, give each level a clear name, and consider a class if the decorator has a lot of state.

  • ↔

    Magic versus explicit calls. A decorator makes behaviour automatic, but a reader must know it is there. For a one-off, calling a helper function explicitly is more honest.

Done when you can

  • Done when you can write a decorator with arguments using three levels.

  • Done when you can explain the error caused by @decorator when @decorator(...) was meant.

  • Done when you can make a decorator that works with and without brackets.

  • Done when you can decorate a method and a class.

  • Done when you can write a decorator as a class with __call__ and update_wrapper.