Декораторы и модуль functools

Одна и та же обвязка повторяется от функции к функции, будь то измерение времени, запись вызова в журнал, проверка аргументов или кеширование результата. Логика везде своя, обвязка одна, а копировать её приходится в каждую функцию. Декораторы позволяют добавить поведение к функции, не затрагивая её тело.

В настоящей главе рассматривается устройство декоратора изнутри, способы передачи ему аргументов, возникающие при этом ловушки, а также готовые декораторы из модуля functools.

Слайды к главе. Материал главы изложен также в пятой и шестой частях лекции «Функции в Python» (декораторы и модуль functools) с демонстрациями в интерактивной оболочке; слайды лекции доступны на сайте книги и в PDF.

Синтаксис декораторов

Декоратор представляет собой вызываемый объект, чаще всего функцию, принимающий другую функцию и возвращающий новую (или любой другой объект). Символ @ перед определением является синтаксическим сахаром, не добавляющим возможностей, но выносящим имя декоратора на видное место перед функцией.

@decorator
def foo(x):
    return 42

Эквивалентно:

def foo(x):
    return 42

foo = decorator(foo)

После применения декоратора имя foo ссылается на результат вызова decorator(foo). Эквивалентность практическая: выражение после @ вычисляется до создания функции, а имя foo с исходной функцией не связывается ни в какой момент. Декоратор выполняется один раз, вместе с инструкцией def, то есть при импорте или запуске модуля, а не при каждом вызове функции. Синтаксис @ появился в Python 2.4 (PEP 318); с версии 3.9 после @ допускается любое выражение, а не только имя или вызов (PEP 614).

В научном коде декораторы встречаются постоянно: @numba.njit компилирует функцию в машинный код, @functools.lru_cache запоминает результаты, @pytest.fixture объявляет подготовку данных для тестов.

Устройство простого декоратора

Рассмотрим устройство на примере декоратора trace, печатающего каждый вызов функции. Декоратор получает исходную функцию в func и определяет новую, inner; та принимает аргументы в самом общем виде *args, **kwargs, применимом к любой функции; inner выполняет свою работу, передаёт управление оригиналу и возвращает полученный результат. Наружу декоратор возвращает inner.

def trace(func):
    def inner(*args, **kwargs):
        print(func.__name__, args, kwargs)
        return func(*args, **kwargs)
    return inner

func остаётся доступной внутри inner даже после завершения trace: это замыкание, рассмотренное в главе про функции.

Применим его.

@trace
def identity(x):
    "I do nothing useful."
    return x

identity(42)  # Вывод: identity (42,) {} → 42

Проблема атрибутов функции

Побочный эффект заключается в том, что после декорирования имя identity ссылается уже не на исходную функцию, а на inner со всеми её атрибутами. Документация, имя и сигнатура берутся теперь от обёртки, и help не покажет ничего полезного.

identity.__name__  # 'inner'
help(identity)     # Справка о inner, а не identity

Помимо неудобства при отладке, это нарушает работу всего, что опирается на имена функций, от автоматически собираемой документации до фреймворков, ищущих обработчики по имени.

Спасение через functools.wraps

functools.wraps является декоратором, копирующим в обёртку метаданные оригинала: __module__, __name__, __qualname__, __doc__ и __annotations__ (с Python 3.12 также __type_params__). Словарь атрибутов оригинала __dict__ дописывается в словарь обёртки, а ссылка на исходную функцию помещается в __wrapped__. При написании собственного декоратора над внутренней функцией всегда ставится @functools.wraps(func).

import functools

def trace(func):
    @functools.wraps(func)
    def inner(*args, **kwargs):
        print(func.__name__, args, kwargs)
        return func(*args, **kwargs)
    return inner

Теперь identity.__name__ и help(identity) возвращают данные оригинала, а inspect.signature(identity), следуя по ссылке __wrapped__, показывает сигнатуру (x) вместо (*args, **kwargs). Та же ссылка позволяет вызвать исходную функцию в обход декоратора, а inspect.unwrap снимает всю цепочку обёрток. Ещё одно следствие касается сериализации: без wraps декорированную функцию нельзя передать модулю pickle, поскольку обёртка называется trace.<locals>.inner, а с wraps можно, что важно для передачи функций в процессы multiprocessing.

Декораторы с аргументами

Иногда декоратор необходимо настроить, например указать trace, куда выводить. Здесь появляется третий уровень вложенности. Запись @trace(sys.stderr) означает два действия: сначала вызывается trace(sys.stderr), и только полученный результат применяется к функции как декоратор, а значит, trace возвращает декоратор, который уже возвращает обёртку.

import sys


def trace(handle):
    def decorator(func):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            print(func.__name__, args, kwargs, file=handle)
            return func(*args, **kwargs)
        return inner
    return decorator

@trace(sys.stderr)
def identity(x):
    return x

Эквивалентно:

decorator = trace(sys.stderr)
identity = decorator(identity)

Декораторы с опциональными аргументами

Два декоратора, приведённых выше, несовместимы по способу применения: один записывается как @trace, другой только как @trace(...). Чтобы работали оба варианта, декоратор анализирует переданное ему значение. Если получена func, декоратор применён без скобок и необходимо сразу возвращать обёртку. Если ничего не получено, вызов выполнен со скобками, и вернуть необходимо декоратор, применяемый следующим шагом.

def trace(func=None, *, handle=sys.stdout):
    if func is None:
        return lambda func: trace(func, handle=handle)
    if not callable(func):
        raise TypeError("настройки trace передаются по имени: trace(handle=...)")

    @functools.wraps(func)
    def inner(*args, **kwargs):
        print(func.__name__, args, kwargs, file=handle)
        return func(*args, **kwargs)
    return inner

Использование:

@trace
def foo(): ...

@trace(handle=sys.stderr)
def bar(): ...

Звёздочка делает handle параметром, передаваемым только по имени, и настройка записывается как @trace(handle=sys.stderr). От ошибочной записи @trace(sys.stderr) звёздочка не защищает: файловый объект, переданный позиционно, всё равно попадает в func. Без проверки декоратор принял бы файл за декорируемую функцию и вернул бы обёртку, а уже при декорировании, когда эта обёртка получит identity вместо аргументов, возникла бы ошибка AttributeError с сообщением, не указывающим на причину. Проверка callable(func) превращает эту ситуацию в понятный TypeError.

Значение sys.stdout по умолчанию вычисляется при выполнении def, поэтому перенаправление contextlib.redirect_stdout, выполненное позже, печать trace не затрагивает: она уходит в исходный поток. Если это важно, по умолчанию указывают handle=None, а sys.stdout подставляют при вызове.

Полезные декораторы на практике

Декоратор @timethis для замера времени

Декоратор для измерения времени выполняет функцию n_iter раз и печатает лучший результат, а не средний. Среднее искажается любой случайной помехой наподобие переключения задач операционной системой, а минимум показывает, на что код способен в отсутствие помех. По той же причине в стандартной библиотеке предусмотрен timeit.repeat, возвращающий список измерений, из которого берётся min. Сам timeit.timeit минимума не вычисляет и возвращает суммарное время всех запусков, которое приходится самостоятельно делить на их число.

import time

def timethis(func=None, *, n_iter=100):
    if func is None:
        return lambda func: timethis(func, n_iter=n_iter)

    @functools.wraps(func)
    def inner(*args, **kwargs):
        print(func.__name__, end=" ... ")
        acc = float("inf")
        for i in range(n_iter):
            tick = time.perf_counter()
            result = func(*args, **kwargs)
            acc = min(acc, time.perf_counter() - tick)
        print(acc)
        return result
    return inner

Декоратор @once для однократного вызова

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

Хранить состояние декоратор мог бы в замыкании, но здесь оно размещено на inner: функция также является объектом и допускает произвольные атрибуты. Флаг при этом виден извне, и some_func.called покажет, вызывалась функция или нет.

def once(func):
    @functools.wraps(func)
    def inner(*args, **kwargs):
        if not inner.called:
            inner.result = func(*args, **kwargs)
            inner.called = True
        return inner.result
    inner.called = False
    return inner

Для функции без аргументов то же поведение даёт готовый @functools.cache (Python 3.9), рассматриваемый ниже.

Декоратор @memoized и мемоизация

Мемоизация представляет собой запоминание результата отдельно для каждого набора аргументов. Функция ведёт словарь «аргументы → результат» и ничего не вычисляет повторно, если ответ уже сохранён. Ключом служит кортеж из позиционных аргументов и отсортированных именованных; сортировка необходима, чтобы f(a=1, b=2) и f(b=2, a=1) попали в одну ячейку.

def memoized(func):
    cache = {}
    mark = object()  # разделитель позиционных и именованных аргументов

    @functools.wraps(func)
    def inner(*args, **kwargs):
        key = args + (mark,) + tuple(sorted(kwargs.items()))
        if key not in cache:
            cache[key] = func(*args, **kwargs)
        return cache[key]
    return inner

Между двумя частями ключа стоит особый объект-разделитель mark. Без него вызовы f(1, ("a", 2)) и f(1, a=2) дали бы одинаковый ключ (1, ("a", 2)) и получили бы один результат на двоих; так же, отдельным маркером, разделяет части ключа и functools.lru_cache.

Ключом словаря может быть только хешируемый объект, поэтому такая мемоизация приведёт к ошибке, если в функцию передать список, словарь или множество. Ради универсальности аргументы сериализуют, например через pickle, но тогда теряется скорость. На практике проще применять мемоизацию только к функциям с простыми аргументами и использовать готовый functools.lru_cache, рассматриваемый ниже.

Декоратор @deprecated для устаревших функций

Когда функцию планируется удалить, её сначала помечают устаревшей. Она ещё работает, но при каждом вызове выдаёт предупреждение; при фильтре по умолчанию оно выводится один раз для каждой строки, из которой функция вызвана. Это даёт пользователям библиотеки время переписать код.

import warnings

def deprecated(func):
    @functools.wraps(func)
    def inner(*args, **kwargs):
        warnings.warn(
            f"{func.__name__} устарела и будет удалена",
            category=DeprecationWarning,
            stacklevel=2,      # показать строку вызывающего кода, а не эту
        )
        return func(*args, **kwargs)
    return inner

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

Аргумент stacklevel=2 переводит указатель со строки внутри декоратора на ту, где функция вызвана, и относит предупреждение к модулю вызывающего кода. Без него место вызова по предупреждению найти невозможно, а при фильтре по умолчанию предупреждение не выводится вовсе.

DeprecationWarning по умолчанию скрыт везде, кроме кода модуля __main__ (PEP 565, Python 3.7). Чтобы увидеть его в собственном коде, интерпретатор запускают с ключом -W default или добавляют строку warnings.simplefilter("always", DeprecationWarning); ключ -W error::DeprecationWarning превращает такие предупреждения в исключения, что удобно в тестах.

С Python 3.13 стандартная библиотека содержит готовый декоратор warnings.deprecated (PEP 702). Он выдаёт то же предупреждение при вызове и вдобавок сообщает об устаревании статическим анализаторам типов.

from warnings import deprecated

@deprecated("old_api устарела, её заменяет new_api")
def old_api(): ...

Контрактное программирование с декораторами

Контрактное программирование состоит в том, чтобы описывать в самой функции условия, требуемые от входа, и обещания, даваемые на выходе. Проверка, добавленная декоратором рядом с объявлением, видна сразу, а тело функции свободно от проверок. Запишем два декоратора, @pre для предусловия и @post для постусловия. Оба принимают проверяющую функцию и текст сообщения: pre проверяет аргументы до вызова, post — результат после.

def pre(cond, message):
    def decorator(func):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            assert cond(*args, **kwargs), message
            return func(*args, **kwargs)
        return inner
    return decorator

def post(cond, message):
    def decorator(func):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            result = func(*args, **kwargs)
            assert cond(result), message
            return result
        return inner
    return decorator

Использование:

import math


@pre(lambda x: 0 < x < 1, "argument must be a fraction")
@post(lambda r: r < 0, "log of a fraction must be negative")
def log_fraction(x):
    return math.log(x)

Предусловие на физический смысл аргумента (доля в интервале от 0 до 1, положительная масса, температура выше абсолютного нуля) останавливает расчёт на входе функции с понятным сообщением, а не порождает nan, обнаруживаемый только в итоговом графике.

assert полностью исключается из кода, если интерпретатор запущен с ключом -O. Для контрактов внутри собственной программы это подходит: при отладке проверка выполняется, а в рабочем режиме не выполняется, и остаётся лишь вызов обёртки. Чтобы убрать и его, декоратор возвращает функцию без обёртки: return inner if __debug__ else func. Данные, поступившие извне, проверять через assert нельзя, там необходим явный raise.

Цепочки декораторов

Декораторов на одной функции может быть несколько, и их порядок важен.

@deco1
@deco2
def foo(): ...

Применяются они снизу вверх, ближайший к def — первым.

foo = deco1(deco2(foo))

Сами выражения после @ вычисляются сверху вниз, что заметно у декораторов с аргументами, а полученные декораторы применяются снизу вверх. При вызове управление передаётся в обратную сторону: сначала во внешний deco1, затем во внутренний deco2 и только потом в саму функцию. Поэтому @functools.lru_cache располагают выше проверок, чтобы попадание в кеш не расходовало время на валидацию, а @timethis располагают ниже, чтобы он измерял саму функцию, а не работу декораторов, применённых сверху.

Декораторы стандартной библиотеки

Большая часть декораторов, встречающихся на практике, уже есть в стандартной библиотеке.

Декоратор применяется и к классу (PEP 3129): он получает класс и возвращает класс. Так устроен @dataclass, дописывающий по аннотациям полей методы __init__, __repr__ и __eq__; классы рассматриваются в главе «Классы».

Модуль functools

Рекурсия и предел глубины

Мемоизация особенно полезна для рекурсивных функций. Каждый рекурсивный вызов создаёт кадр стека со своими аргументами и локальными переменными, и кадры накапливаются, пока рекурсия не дойдёт до базового случая. Интерпретатор ограничивает глубину: sys.getrecursionlimit() по умолчанию возвращает 1000, и более глубокая рекурсия завершается ошибкой RecursionError.

def depth(n):
    return 0 if n == 0 else 1 + depth(n - 1)

depth(900)   # 900
depth(5000)  # RecursionError: maximum recursion depth exceeded

Предел поднимается функцией sys.setrecursionlimit, но это откладывает проблему, а не решает её. Хвостовую рекурсию Python не оптимизирует; это сознательное решение, сохраняющее полную трассировку ошибок. Поэтому глубокая рекурсия переписывается циклом, при необходимости со своим стеком в списке. Предел на практике достигается там, где глубина рекурсии равна размеру данных: рекурсивный поиск связного кластера на решётке (перколяция, кластерный алгоритм Вольфа для модели Изинга) или обход графа в глубину на решётке 100 × 100 легко превышает 1000 уровней, а итеративный обход со стеком из списка такого ограничения не имеет. Сама рекурсия рассмотрена в главе «Рекурсия и сортировки», обход графа в глубину — в главе «Графы».

Другая проблема рекурсии заключается в повторных подзадачах: наивная функция для чисел Фибоначчи вычисляет одни и те же значения экспоненциально много раз.

def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)

Мемоизация с ограничением через lru_cache

Готовая замена рассмотренному выше @memoized. Кеш ограничен по размеру, и при переполнении из него вытесняется запись, к которой дольше всего не обращались, — так расшифровывается LRU, least recently used. Благодаря ограничению кеш не занимает всю память в долго работающей программе.

@functools.lru_cache(maxsize=128)
def expensive_func(x):
    return x * x

С Python 3.8 декоратор записывается и без скобок, @functools.lru_cache, с размером по умолчанию 128. Значение maxsize=None снимает ограничение, и так поступают только для функций с заведомо небольшим числом различных аргументов; с Python 3.9 то же записывается короче: @functools.cache. Метод cache_info() показывает число попаданий и промахов, cache_clear() очищает кеш, а параметр typed=True хранит аргументы разных типов, например 1 и 1.0, раздельно. Если попаданий почти нет, кеш только расходует память.

Ограничение то же, что у самописного варианта: аргументы обязаны быть хешируемыми, поэтому функцию от массива NumPy так не кешируют. В отличие от @memoized, именованные аргументы в ключе не сортируются, и вызовы f(a=1, b=2) и f(b=2, a=1) могут занять две записи кеша.

Кеш превращает экспоненциальную рекурсию для чисел Фибоначчи в линейную, поскольку каждое значение вычисляется однажды. Предел глубины мемоизация при этом не снимает: fib(2000) с пустым кешем завершается RecursionError, поэтому таблицу значений заполняют по возрастанию n или считают циклом.

@functools.cache
def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)

fib(32)           # 2178309
fib.cache_info()  # CacheInfo(hits=30, misses=33, maxsize=None, currsize=33)

Без кеша вычисление fib(32) на Python 3.12 заняло около 0,15 с, с кешем — около 70 мкс.

Частичное применение через partial

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

f = functools.partial(sorted, key=lambda x: x[1])
f([('a', 4), ('b', 2)])  # [('b', 2), ('a', 4)]

int2 = functools.partial(int, base=2)
int2("1010")    # 10
int2.keywords   # {'base': 2}

У partial три преимущества перед лямбдой. Значения аргументов закрепляются в момент создания, поэтому ловушки позднего связывания нет. Объект partial сериализуется модулем pickle, если сериализуемы исходная функция и закреплённые аргументы, и потому передаётся в процессы через multiprocessing.Pool.map, где лямбда вызывает ошибку сериализации; исходная функция при этом определяется в импортируемом модуле, а не в ячейке блокнота (см. главу «Многопоточность и GIL»). Атрибуты func, args и keywords показывают, что именно закреплено. Интегратор quad ожидает функцию одной переменной, а функция Планка зависит от частоты и температуры: partial(planck, T=5800.0) закрепляет температуру.

С Python 3.14 объект functools.Placeholder позволяет закрепить не только первые позиционные аргументы: на его месте остаётся пропуск, заполняемый при вызове.

Обобщённые функции и singledispatch

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

@functools.singledispatch
def pack(obj):
    raise TypeError(f"Unsupported type: {type(obj)}")

@pack.register(int)
def _(obj):
    return b"I" + hex(obj).encode("ascii")

@pack.register(list)
def _(obj):
    return b"L" + b",".join(map(pack, obj))

Выбор реализации выполняется по типу первого аргумента, отсюда «single» в названии, причём с учётом наследования: для True выбирается реализация для int, поскольку bool является подклассом int. Зарегистрированные функции названы _, поскольку обращаться к ним напрямую не требуется: вызываться всегда будет базовая pack.

С Python 3.7 тип указывается и аннотацией первого параметра регистрируемой функции, без аргумента у register; с Python 3.11 в такой аннотации допускается объединение типов. Для методов классов предусмотрен functools.singledispatchmethod (Python 3.8).

@pack.register
def _(obj: str):
    return b"S" + obj.encode("utf-8")

Свёртка последовательности через reduce

reduce сворачивает последовательность в одно значение, беря первые два элемента, применяя к ним функцию, а к полученному результату применяя её же со следующим элементом, и так до конца. Необязательный третий аргумент задаёт начальное значение, которое ставится перед элементами и возвращается для пустой последовательности; без него свёртка пустой последовательности завершается ошибкой TypeError.

functools.reduce(lambda acc, x: acc * x, [1, 2, 3, 4])  # 24

В функциональных языках это базовая конструкция, а в Python она применяется редко: явный цикл читается лучше, а для наиболее частых свёрток предусмотрены готовые sum, min, max, any, all и появившаяся в Python 3.8 math.prod, для которой math.prod([1, 2, 3, 4]) даёт те же 24. Встроенная sum для чисел с плавающей точкой с Python 3.12 выполняет компенсированное суммирование и точнее свёртки через сложение: sum([0.1] * 10) даёт 1.0, а functools.reduce(operator.add, [0.1] * 10) — 0.9999999999999999.

Другие средства functools

Остальные средства модуля относятся к классам, обёрткам и сортировке. total_ordering достраивает все операции сравнения класса по __eq__ и одному из методов __lt__, __le__, __gt__, __ge__; cached_property (Python 3.8) вычисляет атрибут при первом обращении и запоминает значение; partialmethod и singledispatchmethod являются вариантами partial и singledispatch для методов; cmp_to_key переводит функцию сравнения старого стиля в функцию-ключ.

Заключение

Декоратор добавляет функции поведение, не затрагивая её тело, и тем самым не позволяет одной и той же обвязке распространяться по всей программе.

Таким образом, на практике потребуется следующее:

• Собственный декоратор всегда начинается с @functools.wraps(func) над внутренней функцией.
• Декоратор выполняется один раз, вместе с инструкцией def, а обёртка — при каждом вызове функции.
• Декоратор с аргументами представляет собой функцию, возвращающую декоратор; уровней вложенности становится три.
• Декоратор с необязательными настройками принимает их только по имени и проверяет, что первый аргумент является вызываемым объектом: звёздочка сама по себе не защищает от записи @deco(x).
• В цепочке декораторов ближайший к def применяется первым, а при вызове срабатывает последним.
• Прежде чем писать собственную реализацию, следует обратиться к functools: cache, lru_cache, partial, singledispatch и total_ordering покрывают большую часть потребностей.

Дополнительные материалы:

• Python Decorator Library
• PEP 443 – Single-dispatch generic functions
• PEP 318 – Decorators for Functions and Methods
• PEP 702 – Marking deprecations using the type system