Пояснение к лекции про декораторы в CS Центре

with-arguments.md

Хочу ещё раз на примере декоратора trace пояснить, какие типы декораторов используются на практике и как они работают.

Декораторы без аргументов

Общая структура декоратора и пример использования:

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


@trace
def identity(x):
    return x

Применение декоратора trace заменяет имя identity в текущей области видимости на результат вызова trace c текущим значением identity в качестве аргумента:

identity = trace(identity)

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

Решение "в лоб"

Общая структура и пример использования:

def trace(handle):
    def decorator(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

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

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

Декоратор with_arguments

Тройная вложенность версии декоратора trace с аргументами несколько удручает. Но выход есть! Можно обобщить логику декораторов с аргументами в виде декоратора with_arguments.

def with_arguments(deco):
    @functools.wraps(deco)
    def wrapper(*dargs, **dkwargs):                 # 1.
        def decorator(func):                        # 2.
            result = deco(func, *dargs, **dkwargs)  # 3.
            functools.update_wrapper(result, func)
            return result                           # 4.
        return decorator
    return wrapper

Теперь декоратор trace можно переписать так:

@with_arguments
def trace(func, handle):
    # Обратите внимание, что вызывать `functools.wraps` не нужно:
    # это уже делает за нас декоратор `with_arguments`.
    def inner(*args, **kwargs):
        print(func.__name__, args, kwargs, file=handle)
        return func(*args, **kwargs)
    return inner

Разберём, как это работает, по шагам:

1. Имя trace заменяется на результат применения декоратора with_arguments к декоратору trace:

   trace = with_arguments(trace)

   Таким образом, по имени trace теперь доступна функция wrapper, в которой аргумент deco указывает на тело декоратора trace:

   def wrapper(*dargs, **dkwargs):
      def decorator(func):
          result = deco(func, *dargs, **dkwargs)
          functools.update_wrapper(result, func)
          return result
      return decorator

2. Следующий шаг рассмотрим на примере:

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

   Как и в предыдщуей версии, сначала вычисляется выражение после символа @:

   trace_stderr = trace(sys.stderr)

   В результате вызова trace по имени trace_stderr будет доступна функция decorator, в которой аргумент dargs замкнут на значение (sys.stderr, ).

3. Полученный декоратор trace_stderr применяется к функции identity:

   identity = trace_stderr(identity)

   В этот момент вычисляется тело функции decorator. Напомню, что deco указывает на тело декоратора trace, а dargs содержат sys.stderr.

4. В завершении по имени identity записывается значение result из тела функции decorator.

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

Наивная версия

Декоратор with_arguments допускает указание ключевых аргументов. Попробуем это на примере trace:

@with_arguments
def trace(func, handle=sys.stdout):
    def inner(*args, **kwargs):
        print(func.__name__, args, kwargs, file=handle)
        return func(*args, **kwargs)
    return inner

Получившийся декоратор trace можно вызывать без аргументов, но при этом обязательно использовать скобки:

@trace()
def identity(x):
    return x

В противном случае имя identity будет указывать на функцию decorator из тела декоратора with_arguments:

>>> @trace
... def identity(x):
...     return x
...
>>> identity
<function __main__.with_arguments.<locals>.wrapper.<locals>.decorator>

Версия с только ключевыми аргументами

Уйти от лишних скобок можно с помощью только ключевых аргументов:

def trace(func=None, *, handle=sys.stdout):
    # со скобками
    if func is None:
        def decorator(func):
            return trace(func, handle=handle)
        return decorator

    # без скобок
    @functools.wraps(func)
    def inner(*args, **kwargs):
        print(func.__name__, args, kwargs, file=handle)
        return func(*args, **kwargs)
    return inner

Почему это работает?

1. Когда мы вызываем декоратор без скобок, func указывает на декорируемую функцию, а для аругмента handle используется значение по умолчанию.

   @trace
   def identity(x):
       return x

2. Для понимания вызова с handle полезно вспомнить про два шага применения декораторов с аргументами:

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

   На первом шаге trace_stderr = trace(handle=sys.stderr). В этом случае срабатывает ветка func is None и по имени trace_stderr записывается локальная функция decorator, которая фиксирует аргумент handle в значение sys.stderr. На втором шаге identity = trace_stderr(identity), что эквивалентно identity = trace(identity, handle=sys.stderr).