Декорирование функций - это, наверное, самая сложная среди базовых и самая простая среди продвинутых фич языка Python. С декораторами, наверное, знакомы все джуны (хотя бы в рамках подготовки к собеседованиям). Однако, крайне мало разработчиков пишут их правильно. Особенно принимая во внимания тенденции последних нескольких лет к аннотированию всего и вся. Даже популярные open-source проекты (если основная часть их кода была написана до 2018 года) вряд ли дадут вам примеры декораторов, отвечающих всем современным требованиям к коду.

Изначально статья должна была получиться на 2 минуты, но с высоким порогом входа, однако, я не могу просто вывалить контент для 5ти таких же отбитых энтузиастов. Придется пошагово объяснять его всем. Прошу прощения за лонгрид и большое количество кода.

В рамках статьи мы разберемся с декорированием функций в Python от простого к самому сложному. Рассмотрим, как их правильно писать и аннотировать, чтобы другие потребители вашего кода не страдали от близкого знакомства с ним. Уверен, что даже если вы чрезвычайно опытный разработчик, вы найдете для себя полезные советы (хотя и можете пропустить солидную часть материала).

Материал актуален для версий Python3.7-3.11 (и частично 3.12), однако, концептуальные изменения с выходом новых версий могут быть разве что в появлении новых более удобных типов для аннотаций.

Давайте разбираться с декораторами.

И небольшое оглавление, для тех, кто не хочет заблудиться:

• Декораторы

• Простейший декоратор

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

• Декоратор Шредингера

• Аннотирование декораторов

  • typing.Callable

  • typing.TypeVar

  • typing.ParamSpec

  • Аннотируем Шредингера

• Class-based декораторы

• Самое интересное

  • Обозначаем проблему

  • Предлагаем решение

  • Делаем интереснее

  • Последний совет

• Заключение

Кто такие декораторы и зачем они нужны

Декораторы функций - это простой способ модификации поведения любой функции без внесения изменения в ее код. Они часто лежат в основе многих библиотек и вы наверняка с ними знакомы, даже если и никогда не писали (привет FastAPI, Flask и миллионы других).

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

Примеры отличного функционала для реализации через декорирование:

• ретраи - tenacity

• логирование ошибок - loguru

• сериализация входящих данных - FastDepends

• регистрация обработчиков - любой HTTP (и не только) фреймворк

Для простоты понимаю декораторов новичками я всегда вдалбливаю им следующее определение:

Декоратор - это функция, принимающая на вход функцию, и возвращающая (другую) функцию

P.S. На самом деле, это может быть и не функция, а любой Callable объект, да и декорировать классы тоже можно. Но обо всем по порядку.

Простейший декоратор

Итак, давайте напишем самый-самый простой декоратор исходя из его определения.

def decorator(     # это функция
    call           # принимает на вход функцию
):
    def wrapper(*args, **kwargs):
        # код до оригинальной функции
        r = call(*args, **kwargs)
        # код после оригинальной функции
        return r
    return wrapper # возвращает другую функцию

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

Процесс декорирования выглядит просто как вызов функции декоратора:

def call(a: int) -> str:
    return str(a)

decorated_call = decorator(call)

# вызов как обычной функции
assert decorated_call(1) == "1"

Однако, в Python есть немного сахара и на этот случай. Я думаю, следующий синтаксис вам знаком:

@decorator
def call(a: int) -> str:
    return str(a)

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

• ваша оригинальная функция call замещается функцией wrapper, и вы больше не можете получить доступ к call в глобальной видимости;

• в ваш декоратор передается ровно один аргумент - декорируемая функция. Если вы хотите передавать дополнительные аргументы, нужно использовать другие приемы.

Небольшой совет: возьмите в привычку всегда декорировать ваш wrapper с помощью functools.wraps. Эта функция копирует всю служебную метаинформацию о декорируемой функции в функцию-декоратор (название функции, докстринги, список входящих аргументов, их типы и тд).

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

Всего пара "лишних" строк кода, не ленитесь:

from functools import wraps

def decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    return wrapper

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

Если вы более опытный разработчик, то должны были сталкиваться и с подобным синтаксисом (хотя бы в примере выше):

@some_decorator(arg1, arg2)
def func():
    ...

Первая попатка осмысления может вогнать в ступор: декоратор же принимает на вход функцию, а тут какие-то аргументы?

Все достаточно прозаично: some_decorator в этом случае не декоратор, а функция, возвращающая декоратор.

Т.е. процесс декорирования (без сахара) на самом деле выглядит немного по другому:

# обычный декоратор
decorated_func = decorator(call)

# декоратор с параметрами
decorator_wrapper = decorator(arg1, arg2)
decorated_func = decorator_wrapper(call)

# или в одну строку
decorated_func = decorator(arg1, arg2)(call)

Надеюсь, вы уже догадались, как пишется подобный декоратор, но я все-таки приведу пример:

def decorator_wrapper(arg1, arg2):
    def real_decorator(func): # объявляем декоратор
        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper

    return real_decorator     # возвращаем декоратор


@decorator_wrapper(1, 2)
def func():
    ...

Так как декоратор в этом случае - замыкание, то вы имеете доступ к arg1 и arg2 на любом уровне вложенности.

Еще одно важное уточнение: процесс декорирования происходит на этапе первого прохождения интерпретатором Python вашего кода (когда Python получает представление об объектах, которыми он может оперировать в рантайме).

Т.е. в самом рантайме (процесс вызова функции) во всех сценариях декорирования будет выполнен только код, находящийся внутри wrapper'а.

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

Декоратор Шредингера

Так, здесь у нас собрались уже матерые разработчики, поэтому и обсудим юзкейсы по-интереснее.

Вы, конечно, знакомы с декораторами Шредингера, которые можно использовать и так, и так:

@pytest.fixture
def simple_fixture(): ...

@pytest.fixture(scope="session")
def session_fixture(): ...

@dataclass
class SimpleDataclass: ...

@dataclass(slots=True)
class SimpleDataclass: ...

Вы когда-нибудь задумывались как такое написать?

Я вот не задумывался. А потом как задумался! - и сразу попал в ступор.

Ведь это одна функция (точно одна, в питоне же нет перегрузок, или есть), хотя применяется она двумя совершенно разными способами. Объединить их в один тоже как-то проблематично. В первом случае у нас 2 уровня вложенности внутри декоратора, во втором - 3. В общем, загадка.

На самом деле, никакой магии здесь тоже нет. Проблема решается добавлением четвертого уровня вложенности (а то 3ех нам "мало").

таким образом в ваш декоратор передается ровно один аргумент - декорируемая функция

Помните такое? Это наш ключ к решению загадки декоратора Шредингера.

Совмещаем обычный декоратор и декоратор с параметрами вместе и получаем что-то такое:

def schrodinger_decorator(
    call = None,
    *,
    arg1 = None,
    arg2 = None,
): ...

При этом все аргументы должны иметь значения по умолчанию так как в одном случае у нас будет передаваться только call, в другом - только дополнительные аргументы.

Затем мы отдельно пишем декоратор с параметрами еще раз:

def decorator_wrapper(arg1, arg2):
    def real_decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper
    return real_decorator

Теперь нам осталось только разрешить два наших сценария проверкой if call is None.

def schrodinger_decorator(
    call = None,
    *,
    arg1 = None,
    arg2 = None,
):
    wrap_decorator = decorator_wrapper(arg1, arg2)

    # если мы использовали декоратор
    # как декоратор с параметрами
    if call is None:
        return wrap_decorator
    
    # если мы использовали декоратор как обычный
    ## arg1 и arg2 при этом принимают
    ## значения по умолчанию
    else:
        return wrap_decorator(Call)

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

# call=func, arg1=None, arg2=None
@schrodinger_decorator
def func(): ...

# call=None, arg1=1, arg2=2
@schrodinger_decorator(arg1=1, arg2=2)
def func(): ...

Аннотирование декораторов

Современный синтаксис Python невозможно представить без аннотации типов. Аннотация обеспечивает работу автодополнения, а других разработчиков - информацией о типах входящих и исходящих параметров, да и вообще дает много других плюшек, значительно повышающих Developer Experience.

Однако, балуясь с декораторами, эти аннотации очень легко сломать. Поэтому давайте разбираться, как не прострелить себе и другим ноги.

Нам понадобяться следующие братья:

• typing.Callable - базовый тип любого вызываемого объекта в Python (функции в том числе)

• typing.TypeVar - переменная типа, позволяет синхронизировать аннотации различных параметров функции друг с другом

• typing.ParamSpec (до py3.10 использовать из typing_extensions) - общая аннотация произвольных аргументов функции

Достаточно тяжело объяснить без примеров, поэтому давайте на них посмотрим.

typing.Callable

Callable используется для "вызываемых" объектов (функций). Это generic объект, у которого нужно указать 2 типа:

• тип входящих аргументы

• тип результата

Пара примеров:

# любые входящие аргументы, любой результат
Callable[..., Any]

# функция вида `func(a: int, b: float)`
Callable[[int, float], Any]

# функция вида `func(a: int, b: float) -> int`
Callable[[int, float], int]

typing.TypeVar

T = TypeVar("T")

def func(a: T) -> T:
    ...

В данном случае мы объявляем функцию, которая принимает аргумент a любого типа, но ее ответ должен быть того же типа, что a:

int -> int, str -> str, float -> float, e.t.c

Таким образом мы синхронизируем типы в рамках аннотации.

typing.ParamSpec

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

F_Spec = ParamSpec("F_Spec")
F_Return = TypeVar("F_Return")

def decorator(
    call: Callable[
        F_Spec,  # функция с произвольными входными аргументами
        F_Return
    ]
) -> Callable[
    F_Spec,      # функция с теми же входными аргументами
    F_Return
]:
    @wraps(func)
    def wrapper(
        *args: F_Spec.args,      # эти аргументы
        **kwargs: F_Spec.kwargs  # эти аргументы
    ) -> F_Return:
        return call(*args, **kwargs):
    return wrapper

Так выглядит абсолютно правильное аннотирование декораторов с выхода py3.10

И закрепим на примере декоратора с параметрами:

def decorator_wrapper(
    arg1: Any,
    arg2: Any
) -> Callable[  # возвращаем реальный декоратор, который
    # принимает на вход функцию
    [Callable[F_Spec, F_Return]],
    # возвращает такую же функцию
    Callable[F_Spec, F_Return]
]: ...

Кстати, манипулируя типом, возвращаемым декоратором, вы можете изменить автокомплиты задекорированной функциии (но с этим вы поиграетесь сами).

Аннотируем Шредингера

А как нам быть в случае декоратора Шредингера? Он то уж точно ломает стандартную аннотацию функции. При этом тип возвращаемого значения зависит от того, как мы его используем.

Т.е. по факту он имеет 2 возможных варианта декорирования:

# аннотация обычного декоратора
def schrodinger_decorator(
    call: Callable[F_Spec, F_Return],
    *,
    arg1: None = None,
    arg2: None = None,
) -> Callable[F_Spec, F_Return]: ...

# аннотация декоратора с параметрами
def schrodinger_decorator(
    call: None = None,
    *,
    arg1: Any = None,
    arg2: Any = None,
) -> Callable[
    [Callable[F_Spec, F_Return]],
    Callable[F_Spec, F_Return]
]: ...

Хорошо, что есть typing.overload, который создан как раз для этого: данный декоратор позволяет объявить несколько аннотация для одной функции в зависимости от входящих и исходящих аргументов.

Просто объединяем оба варианта в один реальный и добавляем несколько вариантов аннотаций:

# аннотация обычного декоратора
@overload
def schrodinger_decorator(
    call: Callable[F_Spec, F_Return],
    *,
    arg1: None = None,
    arg2: None = None,
) -> Callable[F_Spec, F_Return]:
    # тело должно быть пустым
    # этот код только для аннотации
    ...

# аннотация декоратора с параметрами
@overload
def schrodinger_decorator(
    call: None = None,
    *,
    arg1: Any = None,
    arg2: Any = None,
) -> Callable[
    [Callable[F_Spec, F_Return]],
    Callable[F_Spec, F_Return]
]: ... 

# реальный вариант с объединенной аннотацией
def schrodinger_decorator(
    call: Callable[F_Spec, F_Return] | None = None,
    *,
    arg1: Any | None = None,
    arg2: Any | None = None,
) -> Callable[
    [Callable[F_Spec, F_Return]],
    Callable[F_Spec, F_Return]
] | Callable[
    F_Spec,
    F_Return
]:
    wrap_decorator = decorator_wrapper(arg1, arg2)
    if call is None:
        return wrap_decorator
    else:
        return wrap_decorator(Call)

Class-based декораторы

Как мы все уже точно знаем, процесс декорирования выглядит следующим образом:

decorated_func = decorator(func)
result = decorated_func(*args, **kwargs)

А что если мы хотим сделать что-то такое?

decorated_func = DecoratorClass(func)
result = decorated_func(*args, **kwargs)

Никаких проблем, давайте сделаем!

class DecoratorClass:
    def __init__(self, func):
        # выполняется при декорировании
        self.original_call = func
    
    def __call__(self, *args, **kwargs):
        # выполняется при вызове
        return self.original_call(*args, **kwargs)

А что это нам дает? Ну, а вот это уже самое интересное.

Самое интересное

Обозначаем проблему

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

# любой HTTP фреймворк
@app.get("/")
def handler(...): ...

А что если мы хотим регистрировать одну и ту же функцию несколько раз?

@app.get("/")
@app.get("/index")
def handler(...): ...

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

# пример из FastAPI
class FastAPI:
    def get(route: str):  # декоратор с параметром
        def decorator(func):
            self.routes.append(APIRoute(func))
            return func   # возвращаем без изменений

Теперь, сколько декораторов не вешай, они все будут обрабатывать функцию независимо.

Но что если мы хотим знать, что функции была задекорирована до нас? А кем она была задекорирована? Может быть тогда можно пропустить часть работы? А если мы все-таки хотим добавить фунции дополнительное поведение, но только один раз?

Предлагаем решение

Уже становится сложно. В этом случае нам и поможет class-based декоратор. Мы просто замещаем нашу функцию классом, а во всех следующих декораторах проверяем, пришла нам функция или уже класс-декоратор. На основе этого мы и принимаем решение, что делать дальше.

class FuncWrapper:
    wrapped_call: Callable

    def __init__(self, call):
        if not isinstance(call, FuncWrapper):
            self.wrapped_call = call
        else:
            # функция уже задекорирована
            # обрабатываем этот сценарий

    def __call__(self, *args, **kwargs):
        return self.wrapped_call(*args, **kwargs)


@FuncWrapper
@FuncWrapper
def func(a: int) -> str:
    ...

Однако, нам нужно еще и не сломать аннотирование исходной функции, как мы помним. Тут нам пригодится typing.Generic: этот класс позволяет нам объявлять свой класс как generic и синхронизировать аннотации между различными атрибутами и методами.

from typing import Generic, Callable, ParamSpec, TypeVar

F_Spec = ParamSpec("F_Spec")
F_Return = TypeVar("F_Return")

class FuncWrapper(Generic[F_Spec, F_Return]):
    wrapped_call: Callable[F_Spec, F_Return]

    def __init__(
        self,
        call: Callable[F_Spec, F_Return]
    ) -> None:
        if not isinstance(call, FuncWrapper):
            self.wrapped_call = call

    def __call__(
        self,
        *args: F_Spec.args,
        **kwargs: F_Spec.kwargs,
    ) -> F_Return:
        return self.wrapped_call(*args, **kwargs)

@FuncWrapper
def func(a: int) -> str:
    ...

Однако, почему-то в разных версиях IDE код работает немного по-разному: где-то аннотации корректно переносятся, где-то нет.

Поэтому для себя я использую небольшой трюк, который работает везде: пишем "лишнюю" функцию. Кстати, это же является и решением для написания декораторов с параметрами при таком сценарии.

def decorator(
    call: Callable[F_Spec, F_Return]
) -> FuncWrapper[F_Spec, F_Return]:
    return FuncWrapper(call)

@decorator
def func(a: int) -> str:
    ...

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

class FuncWrapper:
    wrapped_call: Callable
    decorators: list[object]

    def __init__(self, call):
        if not isinstance(call, FuncWrapper):
            self.wrapped_call = call
            self.decorators = []
        else:
            self.wrapped_call = call.wrapped_call
            self.decorators = call.decorators
        self.decorators.append(self)

Делаем интереснее

Хранить информацию о том, кто и как уже задекорировал функцию мы научились. Но проблема в том, что декораторы выше нас знают, кто там был до него. А вот те, кто ниже не знают о декораторах сверху. Непорядок какой-то, не находите?

@FuncWrapper  # знает о соседе снизу
@FuncWrapper  # ничего не знает
def func(a: int) -> str:
    ...

Проблем, вроде бы нет, но они могут быть, если вы имеете доступ к объектам промежуточных декораторов (просто поверьте, моя нога до сих пор болит).

Немного контекста: в новой мажорной версии моего фреймворка для работы с брокерами сообщений Propan вы cможете строить пайплайны обработки данных следующим образом:

@broker.subscriber("in-topic")  # отсюда потребляем сообщения
@broker.publisher("out-topic")  # сюда отправляем ответ
async def handler(msg):
    return "processed"

Соответсвенно, не хорошо заставлять пользователя фреймворка помнить, что publisher должен быть объявлен до subscriber (который использует информацию о зарегистрированных publisher'ах). Пришлось исхитрятся, чтобы они могли быть объявлены в любом порядке.

Решение лежит на уровне метапрограммирования: в случае, если мы декорируем уже задекорированную функцию, мы не копируем поля класса в текущий объект, а действительно используем тот же самый объект. Таким образом, во всех декораторах мы оперируем одним и тем же объектом, и поля, дописанные выше (или даже извне), будут видны на всей глубине "матрешки".

Здесь нам поможет метод __new__.

Метод __new__ - настоящий конструктор класса в Python. Именно он создает объект, поля которого затем инстанциируются в методе __init__. Если мы хотим избежать создания нового объекта при каких-либо условиях, нам нужен именно этот парень.

class FuncWrapper(Generic[F_Spec, F_Return]):
    wrapped_call: Callable[F_Spec, F_Return]
    decorators: list[object]

    def __new__(cls, call: Callable[F_Spec, F_Return]) -> Self:
        if isinstance(call, FuncWrapper):
            # если функция уже задекорирована,
            # возвращаем тот же объект
            return call
        # иначе конструируем новый объект
        return super().__new__(cls)

    def __init__(self, call: Callable[F_Spec, F_Return]) -> None:
        # метод __init__ будет вызван в обоих ветках __new__
        if not isinstance(call, FuncWrapper):
            self.wrapped_call = call
            self.decorators = []
        self.decorators.append(self)


# все все знают
@FuncWrapper
@FuncWrapper
def func(a: int) -> str:
    ...

Последний совет

Некоторые библиотеки определяют, является ли функция асинхронной с помощью стандартного метода asyncio.iscoroutinefunction (тот же FastAPI или мой FastDepends).

# исходный код asyncio
def iscoroutinefunction(func):
    """Return True if func is a decorated coroutine function."""
    return (inspect.iscoroutinefunction(func) or
            getattr(func, '_is_coroutine', None) is _is_coroutine)

Проверка на наличие _is_coroutine является костылем для того, чтобы объект AsyncMock также распознавался как асинхронная функция. Поэтому предпочтительнее использовать именно этот метод, а не inspect.iscoroutinefunction (если у вас возникнет такая потребность). Да, это немного сбивает с толку.

Однако, наш класс-декортатор имеет синхронную реализацию метода __call__:

class FuncWrapper:
    def __call__(self, *args, **kwargs):
        return self.wrapped_call(*args, **kwargs)

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

from asyncio.coroutines import iscoroutinefunction, _is_coroutine

class FuncWrapper(Generic[F_Spec, F_Return]):
    wrapped_call: Callable[F_Spec, F_Return]
    _is_coroutine: None

    def __init__(self, call: Callable[F_Spec, F_Return]) -> None:
        if not isinstance(call, FuncWrapper):
            if iscoroutinefunction(call):
                self._is_coroutine = _is_coroutine
            self.wrapped_call = call

@FuncWrapper
async def func()
    ...

assert iscoroutinefunction(func) is True

Начиная с версии Python3.12 вместо поля _is_coroutine вы можете использовать _is_coroutine_marker, которая позволит модулю inspect понять, что ваша функция асинхронна.

Заключение

Я уверен, что в 99.9% случаев вам не понадобится мое решения с Class-based декоратором с переопределением new и прочей запретной магией. Однако, теперь вы знаете, как покрываются все возможные сценарии использования декоротаров: от самых простых, с исполнением кода до/после декорируемой функции, до самых сложных, когда вам необходимо хранить информацию обо всех примененных к функции декораторах.

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

А в качестве вишенки мы разобрались с тем, как правильно аннотировать декораторы в современном Python. Надеюсь, материал был для вас полезен.