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).