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

Как я сказал ранее, декоратор – по сути функция с аргументом – другой функцией, но как добавить туда еще аргументы, подобно коду ниже?

@lru_cache(maxsize=100)
def sqr(i):
    return i ** 2

Справа от знака собачки (@) в синтаксисе декоратора должен стоять какой-то вызываемый объект (т.е. тот, который можно вызвать как функцию), короче нечто foo, которое будет вызвано, как foo(f) в процессе декорации, где f – декорируемая функция. Под это описание попадают:

• имя функции
• переменная, которой присвоена функция
• экземпляр класса, у которого реализован __call__
• или собственно функциональный вызов func(...), который вернет что-то тоже вызываемое из списка выше

Последний вариант и обеспечивает передачу параметров в декоратор:

@decorator(param=42)
def foo(x, y):
    return x + y

Эквивалентно примерно этому:

foo = decorator(param=42)(foo)

# или

pure_decorator = decorator(param=42)
foo = pure_decorator(foo)

Т. е. сначала вызываем декоратор с параметрами (decorator), он возвращает нам «чистый» декоратор (я назвал его pure_decorator) – функцию с одним параметром, а потом он уже оборачивает исходную функцию foo. Короче функция, которая возвращает декоратор, который возвращаем обернутую функцию. Как же это реализовать? Попробуем на примере декоратор, который повторяет функцию n раз. Начнем с фиксированного n = 5:

from functools import wraps

def repeat(f):
    n = 5
    
    @wraps(f)
    def inner(*args, **kwargs):
        for _ in range(n):
            f(*args, **kwargs)
    return inner

Это нас уже не пугает. Поэтому пристроим сверху еще один этаж – функцию, которая будет принимать n как аргумент. Благодаря механизму замыканий параметры будут доступны для использования внутри вложенных функций.

def repeat(n=5):
    def _repeat(f):
        @wraps(f)
        def inner(*args, **kwargs):
            for _ in range(n):
                f(*args, **kwargs)
        return inner
    # не забываем ее вернуть!
    return _repeat

Вот перед нами три уровня вложенности функций. Это может показаться сложновато, но нужно понять, какая функция для чего нужна:

• repeat нужна, чтобы передать параметр n в _repeat, потому что _repeat не может принять ничего кроме единственного аргумента функции f.
• _repeat нужна, чтобы создать симулякр inner, который умеет принимать любые неважно какие аргументы
• inner нужна, чтобы обернуть вызовы к f и транслировать аргументы f.

Тестируем:

@repeat(3)
def foo():
    print('hello')

foo()
# hello
# hello
# hello

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

twice = repeat(2)

@twice
def bar():
    print('bar')
bar() # два раза вызовет

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

Декоратор, который умеет вызываться с параметром и без

Если мы не укажем n, то repeat будет делать 5 повторов по умолчанию. Можно избавить от необходимости писать пустые скобки, а модифицировать код, чтобы стал более гибким и, работало и так, и так:

@repeat(n=5)
def baz(): ...

@repeat
def baz(): ...

Тут нужна хитрость. Добавим в декоратор скрытый параметр _func = None на самом первом месте. Получится так, что если декоратор вызван без скобок, то единственным его параметром будет декорируемая функция (что подпадает под определение чистого декоратора), которая попадет в переменную _func, а иначе она будет равна None. Если же, мы передаем только именованный параметр n=10, то остается _func = None и благодаря этому декоратор понимает вызвали его с параметром или нет:

def repeat(_func=None, *, n=5):
    def _repeat(f):
        @wraps(f)
        def inner(*args, **kwargs):
            for _ in range(n):
                f(*args, **kwargs)
        return inner

    if _func is None:
        # вызов с параметрами как раньше
        return _repeat
    else:
        # вызов без параметра - сделаем доп. вызов сами
        return _repeat(_func)

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

# уже нельзя
@repeat(10)

# нужно или
@repeat(n=10)
# или
@repeat

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

🤩 Специально для канала @pyway. Подписывайтесь на мой канал в Телеграм @pyway! 👈