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

В этой статье вы узнаете, как использовать декораторы с аргументами и что такое фабрики декораторов.

Предположим, у нас есть функция say(), которая выводит какое-то сообщение:

def say(message):
    ''' выводит сообщение
    Аргументы
        message: сообщение, которое нужно вывести
    '''
    print(message)

и мы хотим выполнять эту функцию 5 раз, каждый раз когда вызываем ее. Например, если мы вызовем функцию с таким аргументом say("Привет"), вывод должен быть таким:

Привет
Привет
Привет
Привет
Привет 

Это можно сделать с помощью простого декоратора:

@repeat
def say(message):
    ''' выводит сообщение
    Аргументы
        message: сообщение, которое нужно вывести
    '''
    print(message)

Для этого просто нужно создать этот декоратор. Например, вот так:

def repeat(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        for _ in range(5):
            result = fn(*args, **kwargs)
        return result

    return wrapper    

Вот весь код:

from functools import wraps

def repeat(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        for _ in range(5):
            result = fn(*args, **kwargs)
        return result

    return wrapper


@repeat
def say(message):
    ''' выводит сообщение
    Аргументы
        message: сообщение, которое нужно вывести
    '''
    print(message)

Но что если нужно выполнить функцию say() 10 раз? В этом случае придется менять захардкоженное значение 5 в декораторе repeat. 

Такое решение — негибкое. Например, если вы захотите использовать декоратор repeat для выполнения функции 5 раз, а потом — 10 раз. Декоратор repeat так сделать не позволит.

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

@repeat(5)
def say(message):
    ...

Чтобы определить декоратор repeat, repeat(5) должен возвращать исходный декоратор.

def repeat(times):
    # возвращает оригинальный декоратор repeat

Новая функция repeat возвращает декоратор. Такие функции часто называют фабрикой декораторов.

Следующая функция repeat возвращает декоратор:

def repeat(times):
    ''' вызывает функцию количество раз, равное times '''
    def decorate(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            for _ in range(times):
                result = fn(*args, **kwargs)
            return result
        return wrapper
    return decorate

В этом коде функция decorate() является декоратором. Она эквивалентна оригинальному декоратору repeat.

Обратите внимание, что новая функция repeat() не является декоратором. Это фабрика декораторов, которая возвращает декоратор.

Соберем все вместе:

from functools import wraps


def repeat(times):
    ''' вызывает функцию количество раз, равное times '''
    def decorate(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            for _ in range(times):
                result = fn(*args, **kwargs)
            return result
        return wrapper
    return decorate


@repeat(10)
def say(message):
    ''' выводит сообщение
    Аргументы
        message: сообщение, которое нужно вывести
    '''
    print(message)


say('Привет')