Декораторы в Python

Краткий конспект по книге “Мощный питон” А. Максвелл. Глава 4.

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

Декоратор – это простая функция, принимающая в качестве аргумента другую функцию (без скобок).
Результат работы декоратора – другая функция.

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

Создание декоратора

def some_decorator(func):  # Это декоратор. Он принимает декорируемую функцию как аргумент

    # Это функция-обертка. Она выполняется всякий раз при вызове декорированной функции
    def wrapper(n):  
        return func(n) + 1

    # Декоратор возвращает функцию-обертку
    return wrapper

Основное свойство декоратора

@some_decorator  # Таким образом происходит оборачивание функции в декоратор
def some_funtion(arg):
    pass

Это выражение равносильно следующему:

def some_function(arg):
    pass
some_function = some_decorator(some_function)

То есть декорированная функция по сути является аргументом функции-декоратора.

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

Уточнение: Её можно вызвать напрямую, если обратиться к атрибуту __wrapped__ (при условии использования @wraps). Это бывает полезно при тестировании. foo.__wrapped__(1) — вызовет оригинал без декоратора.

Зачем нужен @wraps?

При создании своего кастомного декоратора полезно использовать встроенный в Python декоратор @wraps.

from functools import wraps

def some_decorator(func): 

# Необходим, чтобы сохранить __name__, __doc__, __annotations__ декорированной функции.
    @wraps(func)

    def wrapper(n):  
        """Это функция-обертка"""
        return func(n) + 1

    return wrapper

Если не добавлять @wraps, то при вызове декорированной функции и ее атрибутов мы увидим информацию о функции wrapper. Это затрудняет отладку кода.

func(1)
print(func.__name__)
print(func.__doc__)
print(func.__annotations__)

Вывод:
wrapper
Это функция-обертка
{}

Универсальный декоратор, принимающий любое число аргументов

def printlog_1(func):
    @wraps(func)
    def wrapper(*args, **kwargs):  # может принимать любое число аргументов
        print("CALLING: " + func.__name__)
        return func(*args, **kwargs)

    return wrapper

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

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

def countcalls(func):
    count = 0  # Это счетчик

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

        # Ключевое слово nonlocal показывает функции wrapper, что переменная count находится на следующем уровне области видимости, но не является глобальной
        nonlocal count  
        count += 1
        print(f"# of calls: {count}")
        return func(*args, **kwargs)

    return wrapper

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

Подробный пример с сохранением данных в словаре data.

"""Некоторые задачи требуют хранить данные в функции-декораторе - стр. 77"""


def running_average(func):
    """Эта функция-декоратор умеет хранить данные многочисленных запусков
    декорированной функции в словаре data"""

    # Эта часть декоратора срабатывает только один раз - при первом вызове
    # декорированной функции. Все остальные вызовы этой функции приводят к
    # к работе только функции wrapper.
    # Именно благодаря этому словарь data накапливает информацию и работает независимо
    # для разных декорированных функций.
    data = {"total": 0, "count": 0}

    @wraps(func)  # Необходим, чтобы сохранить __name__, __doc__, __annotations__ декорированной функции.
    def wrapper(*args, **kwargs):
        val = func(*args, **kwargs)
        data["total"] += val
        data["count"] += 1

        # В реальном проекте эти данные удобнее записывать в лог, а не выводить на печать
        print("Average of {} so far: {:.01f}".format(
            func.__name__, data["total"] / data["count"]
        ))
        return val

    # Здесь словарь data присваивается функции wrapper в качестве атрибута, и теперь мы можем к нему обращаться извне.
    # Например, из декорированной функции - foo.data
    # Это один из приемов аннотирования функций - прикрепления к ним дополнительной информации.
    wrapper.data = data

    return wrapper


@running_average
def foo(x: int):
    """Добавляет 2 к аргументу"""
    return x + 2


@running_average
def buzz(x, y):
    return x ** y


foo(1)
print(foo.__name__)
print(foo.__doc__)
print(foo.__annotations__)
print(foo.data)
print()
foo(10)
print(foo.data)
foo(1)
print(foo.data)
foo(1)
print(foo.data)

buzz(2, 3)
print(buzz.data)
buzz(3, 4)
print(buzz.data)
buzz(5, 6)
print(buzz.data)

Вывод:
Здесь хорошо видно, что данные для двух разных функций накапливаются независимо.

Average of foo so far: 3.0
wrapper
None
{}
{'total': 3, 'count': 1}

Average of foo so far: 7.5
{'total': 15, 'count': 2}
Average of foo so far: 6.0
{'total': 18, 'count': 3}
Average of foo so far: 5.2
{'total': 21, 'count': 4}
Average of buzz so far: 8.0
{'total': 8, 'count': 1}
Average of buzz so far: 44.5
{'total': 89, 'count': 2}
Average of buzz so far: 5238.0
{'total': 15714, 'count': 3}

Как передать аргумент в декоратор?

"""Как передать аргумент в декоратор?
Необходимо обернуть уже имеющийся декоратор еще одной функцией,
которая принимает нужный аргумент."""


def add(increment):
    def decorator(func):
        @wraps(func)
        def wrapper(n):
            return func(n) + increment

        return wrapper

    return decorator


# add = add(20)  # Можно сначала присвоить декоратор с аргументом отдельному объекту
# @add           # И далее вызвать его перед функцией
# def foo(x):
#     return x + 10
# print(foo(5))

@add(20)  # Но можно и так
def foo(x):
    return x + 10


print(foo(5))

Вывод:
35

Как создать декоратор в виде класса?

"""Декоратор можно реализовать не только в виде функции, но и в виде класса.
Это открывает возможности для использования всех свойств ООП, включая наследование."""


class Prefixer:
    def __init__(self, prefix):
        self.prefix = prefix

    def __call__(self, message):  # Этот метод делает класс вызываемым, подобно функции, принимающей аргумент message
        return self.prefix + message


# Теперь можно вызывать экземпляры этого класса как функции
simonsays = Prefixer("Simon says: ")  # Создаем экземпляр
print(simonsays("Get up and dance!!!"))  # Вызываем как функцию


# Превратим декоратор printlog из функции в класс
# Так он выглядел сначала
def printlog(func):
    def wrapper(*args, **kwargs):  # может принимать любое число аргументов
        print("CALLING: " + func.__name__)
        return func(*args, **kwargs)

    return wrapper


# И так применялся к функции
@printlog
def buzz(x, y):
    print(x + y)


buzz(10, 20)


# Перепишем декоратор в виде класса
class PrintLog:
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        print(f"CALLING: {self.func.__name__}")
        return self.func(*args, **kwargs)


@PrintLog
def buzz(x, y):
    print(x + y)


buzz(50, 60)