PythonMiddleLive coding

Как написать декоратор, который принимает аргументы?

Это трёхуровневая конструкция: внешняя функция-фабрика принимает параметры и возвращает настоящий декоратор; декоратор принимает функцию и возвращает wrapper. @retry(times=3) сначала вызывает retry(times=3), результат применяется к функции.

Три уровня

Декоратор с аргументами — это не декоратор, а фабрика декораторов. Когда вы пишете @retry(times=3), Python сначала вычисляет retry(times=3) (получает реальный декоратор), а потом применяет его к функции:

@retry(times=3)
def f(): ...

# эквивалентно
decorator = retry(times=3)
f = decorator(f)

Готовый пример

from functools import wraps
import time, logging

log = logging.getLogger(__name__)

def retry(times: int = 3, delay: float = 0.5, exceptions: tuple = (Exception,)):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exc = None
            for attempt in range(1, times + 1):
                try:
                    return func(*args, **kwargs)
                except exceptions as exc:
                    last_exc = exc
                    log.warning("attempt %d/%d failed: %s", attempt, times, exc)
                    time.sleep(delay * attempt)
            raise last_exc
        return wrapper
    return decorator

@retry(times=3, delay=0.2, exceptions=(ConnectionError,))
def fetch(url): ...

Поддержка обоих вариантов (@deco и @deco(...))

Иногда хочется, чтобы декоратор работал и без скобок, и со скобками. Шаблон:

from functools import wraps

def cache(func=None, *, maxsize: int = 128):
    if func is None:
        # вызвали @cache(maxsize=...)
        def decorator(f):
            return cache(f, maxsize=maxsize)
        return decorator
    # вызвали @cache без скобок
    storage = {}
    @wraps(func)
    def wrapper(*args, **kwargs):
        key = (args, tuple(sorted(kwargs.items())))
        if key not in storage:
            if len(storage) >= maxsize:
                storage.pop(next(iter(storage)))
            storage[key] = func(*args, **kwargs)
        return storage[key]
    return wrapper

@cache                         # без скобок
def f(x): return x * 2

@cache(maxsize=256)            # со скобками
def g(x): return x * 2

Класс как фабрика

class Retry:
    def __init__(self, times: int = 3):
        self.times = times

    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for _ in range(self.times):
                try:
                    return func(*args, **kwargs)
                except Exception:
                    continue
            raise
        return wrapper

@Retry(times=5)
def get(): ...

Подводные камни

  • Забыть скобки: @retry вместо @retry() — Python вызовет retry(f), передав функцию вместо параметров.
  • Параметры фабрики и аргументы wrapper живут в разных замыканиях — не путайте.
  • Mutable default argument на фабрике (exceptions=[Exception]) — те же грабли, что и у обычных функций.
  • Параметры — позиционные, а вызывающий передал именованные — TypeError в момент импорта, что трудно поймать на этапе CI.
  • Без @wraps метаданные исходной функции теряются.

Common mistakes

  • Писать только двухуровневый декоратор.
  • Возвращать wrapper из внешней функции напрямую.
  • Не понимать порядок вызовов при импорте.

What the interviewer is testing

  • Объясняет три уровня функций.
  • Показывает замыкание с параметром.
  • Знает эквивалент decorator = repeat(3).

Sources

Related topics