Python: декоратор с параметрами и functools.wraps

Категория: Python · Язык: Python · Обновлено:

Декоратор в Python это функция, которая принимает другую функцию и возвращает её изменённую версию. Если декоратору нужны параметры, как в @repeat(3), добавляется ещё один уровень вложенности: внешняя функция принимает параметры, средняя функцию, внутренняя wrapper выполняет работу.

Простой декоратор и functools.wraps

Python
from functools import wraps

def log_call(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(f"вызов {func.__name__}{args}")
        return func(*args, **kwargs)
    return wrapper

@log_call
def add(a, b):
    """Складывает два числа."""
    return a + b

result = add(2, 3)       # печатает: вызов add(2, 3)
print(result)            # 5
print(add.__name__)      # add
print(add.__doc__)       # Складывает два числа.

Запись @log_call перед функцией равна add = log_call(add). Без @wraps(func) имя и строка документации заменились бы на wrapper, а это ломает отладку, справку и инструменты.

Декоратор с параметрами

Python
def repeat(times):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            result = None
            for _ in range(times):
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator

counter = {"n": 0}

@repeat(3)
def hit():
    counter["n"] += 1

hit()
print(counter["n"])      # 3

Запись @repeat(3) выполняется в два шага: сначала вызывается repeat(3) и возвращает decorator, затем decorator применяется к функции.

Повтор при ошибке (retry)

Python
def retry(attempts=3, exceptions=(ValueError,)):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, attempts + 1):
                try:
                    return func(*args, **kwargs)
                except exceptions:
                    if attempt == attempts:
                        raise
        return wrapper
    return decorator

calls = []

@retry(attempts=3)
def flaky():
    calls.append(1)
    if len(calls) < 3:
        raise ValueError("сбой")
    return "ок"

print(flaky())           # ок
print(len(calls))        # 3

Работает с параметрами и без них

Чтобы использовать и @tag, и @tag(name="x"), принимайте функцию необязательным аргументом:

Python
def tag(func=None, *, name="log"):
    def decorate(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            return f"[{name}] {f(*args, **kwargs)}"
        return wrapper
    return decorate(func) if func else decorate

@tag
def first():
    return "один"

@tag(name="new")
def second():
    return "два"

print(first())           # [log] один
print(second())          # [new] два

Несколько декораторов

Запись

Эквивалент

@a над @b над def f

f = a(b(f))

Применяются снизу вверх, а выполняются сверху вниз: сначала внешний wrapper от a, затем от b.

Частые вопросы

Зачем нужен functools.wraps?

Он копирует имя, документацию и атрибуты исходной функции в wrapper. Без него help(add) и трассировки покажут имя wrapper.

Как декорировать метод класса?

Так же. Первым аргументом в *args придёт self, поэтому wrapper должен принимать *args и **kwargs, а не фиксированный набор.

Как узнать, что функция задекорирована?

У результата wraps есть атрибут __wrapped__, он ведёт к исходной функции: add.__wrapped__(2, 3).