Декораторы

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

Декорация функции одного аргумента

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

import math 

def sin(x):
    print(f"Функция sin вызвана с аргументом {x}")
    return math.sin(x)

print(sin(0))
Функция sin вызвана с аргументом 0
0.0

Такие конструкции называют обертками над функциями: пользовательская функция sin оборачивает функцию math.sin.

Теперь, предположим, что мы захотим проделать то же самое и с другими функциями. Написание своей обертки для каждой из них приведет к излишнему дублированию кода. Чтобы этого избежать, напишем функцию announce, которая, принимая на вход некоторую функцию одного аргумента func, возвращает функцию обертку wrap. Функция обертка wrap в свою очередь вызывает func, но перед этим печатает её имя (func.__name__) и значение аргумента, с которым она была вызвана.

Реализация этой функции может выглядеть как-то так.

def announce(func):
    def wrap(x):
        print(f"Функция {func.__name__} вызвана с аргументом {x}")
        return func(x)
    return wrap

Функция announce внутри тела объявляет новую функцию wrap и тем самым создаёт объект этой функции. Когда этот объект создан, он сразу возвращается наружу. Функция wrap в свою очередь является оберткой: она возвращает результат вычисления обертываемой функции func, но перед этим делает дополнительные действия, а именно печатает имя функции (атрибут func.__name__) и аргумент, с которым вызывается функция. С этим же аргументом вызывается функция func.

Проверим работоспособность это функции. Для этого обернем функции синуса и косинуса.

from math import sin, cos

print(sin(0), cos(0))

sin = announce(sin)
cos = announce(cos)

print(sin(0), cos(0))
0.0 1.0
Функция sin вызвана с аргументом 0
Функция cos вызвана с аргументом 0
0.0 1.0

Видно, что обертывание функций sin и cos приводит к желаемому эффекту.

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

Вместо комбинации инструкций

def my_function(x):
    pass

my_function = announce(my_function)

можно смело писать

@announce
def my_function(x):
    pass

Они эквиваленты между собой с точки зрения синтаксиса.

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

@announce
def square(x):
    return x * x

print(square(2))
Функция square вызвана с аргументом 2
4

Проблема с затиранием имени функции и её решение

Попробуем получить имя декорированной функции square из предыдущего примера.

square.__name__
'wrap'

Упс. Вместо ожидаемого square мы получили wrap. Это объясняется тем, что

@announce
def square(x):

сводится к

square = announce(square)

А функция announce возвращает обертку wrap. Таким образом, имя исходной функции затирается.

Можно попробовать решить эту проблему задавав необходимое имя функции wrap перед её возвращением из функции announce.

def announce(func):
    def wrap(x):
        print(f"Функция {func.__name__} вызвана с аргументом {x}")
        return func(x)
    wrap.__name__ = func.__name__
    return wrap

@announce
def square(x):
    return x * x

print(square.__name__)
print(square)
square
<function announce.<locals>.wrap at 0x000001556722AB80>

Но куда лучше будет воспользоваться декоратором wraps из модуля functools, что позволит сохранить все метаданные исходной функции.

from functools import wraps

def announce(func):
    @wraps(func)
    def wrap(x):
        print(f"Функция {func.__name__} вызвана с аргументом {x}")
        return func(x)
    return wrap

@announce
def square(x):
    return x * x

print(square.__name__)
print(square)
square
<function square at 0x000001556722A9D0>

Декорация функции с произвольной сигнатурой

В прошлом примере мы наложили ограничение на декорируемую функцию в один аргумент. Это было сделано лишь для наглядности. Используя *args и **kwargs можно объявлять декораторы для функций с произвольной сигнатурой (см. args и kwargs, распаковка кортежей и словарей).

Переопределим декоратор announce.

def announce(func):
    def wrap(*args, **kwargs):
        print(f"Функция {func.__name__} вызвана с позиционными аргументами {args}", end="")
        print(f" и именованными аргументами {kwargs}.")
        return func(*args, **kwargs)
    return wrap

@announce
def add(x, y):
    return x + y

print(add(42, 3.14))
print(add(7, y=2.71))
Функция add вызвана с позиционными аргументами (42, 3.14) и именованными аргументами {}.
45.14
Функция add вызвана с позиционными аргументами (7,) и именованными аргументами {'y': 2.71}.
9.71

Пример: измеряющий время декоратор

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

from time import perf_counter
import seaborn as sns


def timed(func):
    def wrap(*args, **kwargs):
        t1 = perf_counter()
        result = func(*args, **kwargs)
        t2 = perf_counter()
        print(f"Вызов {func.__name__} занял {t2- t1} секунд, ", end="")
        print(f"(параметры {args}, {kwargs})")
        return result
    return wrap


@timed
def countdown(n):
    while n > 0:
        n -= 1


for n in [10_000, 100_000, 1_000_000, 10_000_000]:
    countdown(n)
Вызов countdown занял 0.0006112999999992041 секунд, (параметры (10000,), {})
Вызов countdown занял 0.006321100000000968 секунд, (параметры (100000,), {})
Вызов countdown занял 0.10701859999999996 секунд, (параметры (1000000,), {})
Вызов countdown занял 0.6229589999999998 секунд, (параметры (10000000,), {})

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

На примере декоратора functools.wraps, декораторы, принимающие аргументы, объявляются несколько сложнее. Дело в том, что в этом случае необходимо погрузится ещё на один уровень абстракции. Предположим, что вы хотите написать декоратор my_decorator с параметрами x, y и z, а применять его в следующем виде:

@my_decorator(x, y, z)
def some_function():
    ...

Такое объявление будет эквивалентно инструкции

some_function = my_decorator(x, y, z)(some_function)

Иными словами, вызов функции my_decorator(x, y, z) должен возвращать другую функцию (или вызываемый объект), которая уже и будет параметризованным декоратором. Т.е. выглядеть она должна приблизительно так:

def my_decorator(x, y, z):
    ...
    def decorate(func):
        ...
        @functools.wraps
        def wrap(*args, **kwargs):
            ...
            result = func(*args, **kwargs)
            ...
            return result
        ...
        return wrap
        ...
    return decorate 

Таким образом инструкция

some_function = my_decorator(x, y, z)(some_function)

сводится к инструкции

some_function = decorate(some_function)

где функция decorate — то, что вернул вызов функции my_decorator(x, y, z) — и является непосредственным декоратором. Функция decorate может получить доступ к параметрам x, y и z с которыми была вызвана функция my_decorator, т.к. они находятся в пространстве локальных имен внешней для неё функции my_decorator. Таким образом поведение функции decorate может зависеть от параметров x, y и z.

В качестве примера, рассмотрим два декоратора. Первый декорирует функцию таким образом, чтобы печаталось заданное сообщение перед вызовом функции, а второй — после вызова функции.

from functools import wraps

def print_before(msg=""):
    def decorate(func):
        @wraps(func)
        def wrap(*args, **kwargs):
            print(msg)
            return func(*args, **kwargs)
        return wrap
    return decorate

def print_after(msg):
    def decorate(func):
        @wraps(func)
        def wrap(*args, **kwargs):
            result = func(*args, **kwargs)
            print(msg)
            return result
        return wrap
    return decorate

Приблизительно так может выглядеть реализация таких декораторов.

Проверим работоспособность первого из них.

@print_before(msg="Hello!")
def f():
    print("Функция f.")

f()
Hello!
Функция f.

Проверим работоспособность второго из них.

@print_after(msg="Good bye!")
def g():
    print("Функция g.")

g()
Функция g.
Good bye!

Теперь убедимся, что мы можем передать декоратору и другое значение, а заодно продемонстрируем возможность цепного применения декораторов.

@print_before(msg="Hi!")
@print_after(msg="Bye!")
def h():
    print("Функция h.")

h()
Hi!
Функция h.
Bye!

Note

Объявление декораторов с необязательными аргументами и другие примеры можно найти в [2].