Декораторы в Python: обёртка вокруг функции

Вы уже умеете передавать функцию в функцию и возвращать её из функции — это были функции как значения. Декоратор — это функция, которая принимает другую функцию и возвращает новую: поведение добавляется без правки исходного кода. Логирование, замер времени, проверка доступа — всё это дописывается обёрткой снаружи, а функция внутри остаётся нетронутой. Само слово страшнее механизма: разберём декоратор по косточкам, без знака @.

Обёртка руками: без волшебного знака

Задача: сделать из тихой функции громкую — перевести её ответ в верхний регистр и добавить восклицание. Пишем функцию shout, которая принимает функцию и возвращает обёрнутую:

def shout(func):
    def wrapper():
        result = func()
        return result.upper() + '!'
    return wrapper

def greet():
    return 'привет'

greet = shout(greet)
print(greet())

ПРИВЕТ!

Разбираем: wrapper — это замыкание из прошлого урока, оно помнит исходную func. Строка greet = shout(greet) подменила функцию: имя greet теперь указывает на обёртку. Вызов greet() на самом деле зовёт wrapper, а тот внутри вызывает оригинал и приукрашивает результат. Исходная функция не изменилась ни на символ — поведение добавилось снаружи.

Синтаксис @: та же строка, но короче

Паттерн «объявили функцию, потом подменили её обёрткой» встречается так часто, что для него завели короткую запись — знак @ над объявлением:

def shout(func):
    def wrapper():
        return func().upper() + '!'
    return wrapper

@shout
def greet():
    return 'привет'

@shout
def farewell():
    return 'пока'

print(greet())
print(farewell())

ПРИВЕТ!
ПОКА!

Обёртка, пропускающая аргументы насквозь

Обёртка из примера зовёт func без аргументов, а настоящие функции принимают параметры. Решение — связка *args и **kwargs: обёртка соглашается взять любые аргументы и передаёт их оригиналу в том же виде.

def log_call(func):
    def wrapper(*args, **kwargs):
        print('Вызов:', func.__name__)
        return func(*args, **kwargs)
    return wrapper

@log_call
def add(a, b):
    return a + b

@log_call
def greet(name):
    return f'Привет, {name}!'

print(add(2, 3))
print(greet('Анна'))

Вызов: add
5
Вызов: greet
Привет, Анна!

Один декоратор обслужил функции с разными сигнатурами: и add(2, 3), и greet('Анна') прошли через одну и ту же обёртку. Атрибут func.__name__ — имя исходной функции, обычная строка, по ней обёртка пишет, кого вызвали. Обратите внимание на порядок строк вывода: сначала журнал обёртки, потом результат функции.

Декоратор-счётчик

Обёртка может не только пропускать вызов насквозь, но и вести свою бухгалтерию. Счётчик вызовов приклеим к самой обёртке — функции в Python можно добавлять атрибуты, как переменные:

def counted(func):
    def wrapper(*args):
        wrapper.calls += 1
        print('Вызов номер', wrapper.calls)
        return func(*args)
    wrapper.calls = 0
    return wrapper

@counted
def square(n):
    return n * n

print(square(3))
print(square(4))
print(square(5))
print('Всего вызовов:', square.calls)

Вызов номер 1
9
Вызов номер 2
16
Вызов номер 3
25
Всего вызовов: 3

Строка wrapper.calls = 0 выполняется один раз при создании обёртки, дальше счёт растёт при каждом вызове. Состояние снова спрятано в замыкание — глобальная переменная не понадобилась.

Честное имя: functools.wraps

У обёртки есть цена: подменив функцию, декоратор подменил и её представление. Функция greet теперь втайне зовётся wrapper:

def shout(func):
    def wrapper():
        return func().upper() + '!'
    return wrapper

@shout
def greet():
    return 'привет'

print(greet())
print(greet.__name__)

ПРИВЕТ!
wrapper

Для вызовов это неважно, а вот для отладки и документации — да: стек ошибок покажет wrapper, справка потеряется. Стандартная библиотека лечит это декоратором wraps из модуля functools: он копирует имя и документацию оригинала на обёртку.

from functools import wraps

def shout(func):
    @wraps(func)
    def wrapper():
        return func().upper() + '!'
    return wrapper

@shout
def greet():
    return 'привет'

print(greet())
print(greet.__name__)

ПРИВЕТ!
greet

Шаблон вырос на две строки — from functools import wraps и @wraps(func) над wrapper — но теперь обёртка ведёт себя как вежливый двойник: работает за оригинал и носит его имя. Пишите @wraps в каждом декораторе сразу, привычка дешёвая, а отладка благодарная.

Замер времени: честная оговорка

import time

def timed(func):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        elapsed = time.perf_counter() - start
        print(f'Время: {elapsed:.4f} с')
        return result
    return wrapper

Это классика жанра: журнал и счётчик из урока можно запускать — их вывод всегда одинаков; таймер печатает число, которое каждый раз новое. Правило простое: то, что декоратор печатает, должно быть предсказуемым, иначе примеры нельзя проверить.

Стек декораторов: кто первый

Над одной функцией можно поставить два декоратора. Применяются они снизу вверх: декоратор ближе к def — первым:

def reverse_text(func):
    def wrapper():
        return func()[::-1]
    return wrapper

def mark(func):
    def wrapper():
        return func() + '#'
    return wrapper

@reverse_text
@mark
def first():
    return 'кот'

@mark
@reverse_text
def second():
    return 'кот'

print(first())
print(second())

#ток
ток#

В первом случае mark отработал раньше: метка приклеилась, потом строка перевернулась — решётка оказалась впереди. Во втором — наоборот. Одна и та же пара декораторов, разный порядок — разный результат: читайте стек снизу вверх, как историю выполнения.

Декораторы, которые вы уже встретите

• логирование вызовов: журнал из блока «логирующий декоратор» — кто, когда и с чем;
• счётчики: сколько раз вызвали функцию, как часто стреляет медленный кусок;
• проверки: обёртка может заранее проверить аргументы и не пустить их в функцию;
• кэш результатов: модуль functools содержит готовый декоратор lru_cache — он помнит ответы на уже заданные вопросы.

Декоратор собран из трёх вещей, которые вы уже знаете: функция как значение, замыкание и *args с **kwargs. Осталось привыкнуть к записи @ и к тому, что имя функции после декорации указывает на обёртку. Когда что-то пойдёт не так — обёртка примет не тот аргумент или забытый return wrapper вернёт None — на помощь придёт следующий урок об ошибках вызова.

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

Иногда декоратору нужны свои настройки: повторить вызов три раза, повторять попытку при ошибке, печатать только при verbose. Тогда в записи появляются скобки — @repeat(times=3) — и внутри появляется ещё один уровень вложенности.

from functools import wraps

def repeat(times):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            return [func(*args, **kwargs) for _ in range(times)]
        return wrapper
    return decorator

@repeat(times=3)
def greet(name):
    return f'Привет, {name}!'

print(greet('Анна'))

['Привет, Анна!', 'Привет, Анна!', 'Привет, Анна!']

Три уровня вложенности, и каждый отвечает за своё. repeat — фабрика: она принимает настройку times и возвращает декоратор. decorator — сам декоратор: он принимает функцию и возвращает обёртку. wrapper — обёртка: она выполняет функцию столько раз, сколько велено. Поэтому в вызове и скобки: @repeat(times=3) сначала вызывает repeat и получает декоратор, а тот уже применяется к greet. У @shout скобок не было — фабрика не требовала настроек.

Ограничение одно, и оно важное: декоратор не может изменить тело функции. Он видит только аргументы и результат — всё, что происходит внутри, ему недоступно. Если логике нужен доступ к внутреннему состоянию объекта, это уже не декоратор, а обычный метод.

Декоратор на методе класса

Декораторы применяют и к методам — и там есть одна деталь, на которой спотыкаются: объект self. Спасает всё та же универсальная сигнатура обёртки.

from functools import wraps

def logged(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print('Вызов:', func.__name__)
        return func(*args, **kwargs)
    return wrapper

class Report:
    @logged
    def build(self, name):
        return f'отчёт {name}'

report = Report()
print(report.build('июнь'))

Вызов: build
отчёт июнь

self попал в args вместе с остальными аргументами — обёртка ничего не знает о классах и просто пересылает всё насквозь. Если бы она была написана с фиксированной сигнатурой wrapper(name), то self занял бы место name, и метод сломался бы с невнятной ошибкой. Общее правило для обёрток: принимать *args, **kwargs и не угадывать сигнатуру.

Декоратор-проверка

Ещё один полезный класс декораторов — те, что не дают функции выполниться с плохими данными. Проверка живёт в обёртке, а сама функция остаётся чистой: она получает хорошие значения и работает с ними.

from functools import wraps

def positive(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        for value in args:
            if value <= 0:
                return 'ошибка: нужно положительное число'
        return func(*args, **kwargs)
    return wrapper

@positive
def square(n):
    return n * n

print(square(5))
print(square(-2))

25
ошибка: нужно положительное число

Плохое значение остановилось на обёртке, функция square даже не вызвалась. В рабочем коде вместо возврата строки здесь подняли бы исключение — но для примера мы обошлись строкой, чтобы обе ветки были видны. Такие проверки экономят время: вместо падения в глубине расчёта вы получаете понятное сообщение на входе.