Декораторы

Декоратор это самый частый вопрос по Python на собеседованиях инженеров данных после генераторов, и самый частый повод для ответа «ну, это такая штука с собачкой». Между тем за ним нет ничего, кроме двух фактов из прошлого урока: функция это объект, и внутренняя функция помнит переменные внешней.

Что происходит при @

def logged(func):
    def wrapper(*args, **kwargs):
        print("вызов", func.__name__)
        return func(*args, **kwargs)
    return wrapper
 
@logged
def load(table):
    return f"{table}: ok"

Запись с @ это сокращение для присваивания:

def load(table):
    return f"{table}: ok"
 
load = logged(load)

logged получает функцию, создаёт новую (wrapper), которая делает что-то до и после вызова исходной, и возвращает её. Имя load теперь указывает на wrapper, а исходная функция живёт внутри замыкания как func. Никакой магии: декоратор это функция, принимающая функцию и возвращающая функцию.

*args, **kwargs в обёртке нужны, чтобы она подходила к функции с любой сигнатурой и передавала всё как есть. return func(...) обязателен: без него декорированная функция начнёт возвращать None, и это классическая ошибка в первом самописном декораторе.

functools.wraps

После декорирования load.__name__ равен "wrapper", docstring потерян, в traceback и логах Airflow видно wrapper вместо имени задачи. functools.wraps копирует метаданные исходной функции на обёртку:

import functools
 
def logged(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print("вызов", func.__name__)
        return func(*args, **kwargs)
    return wrapper

Это не опция, а обязательная строка в любом декораторе. Без неё ломаются инструменты, которые смотрят на имя функции: тесты, профилировщики, сам Airflow при регистрации задач.

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

@retry(times=3) это не декоратор, а вызов функции, которая возвращает декоратор. Отсюда три уровня вложенности:

import functools
import time
 
def retry(times, delay=1.0, exceptions=(Exception,)):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            last_error = None
            for attempt in range(times):
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    last_error = e
                    time.sleep(delay * (attempt + 1))
            raise last_error
        return wrapper
    return decorator
 
@retry(times=3, exceptions=(ConnectionError, TimeoutError))
def fetch_page(url):
    ...

retry(times=3) выполняется первым и возвращает decorator, затем decorator(fetch_page) возвращает wrapper. Внешние параметры (times, delay) доступны в wrapper через замыкание.

Обратите внимание на exceptions: повторять стоит только то, что может пройти со второго раза. Сетевая ошибка да, KeyError из-за кривых данных нет, он будет кривым и на третьей попытке. Декоратор, который ловит всё подряд, превращает ошибку данных в трёхкратную задержку и тот же самый traceback.

Несколько декораторов

Декораторы применяются снизу вверх, а выполняются сверху вниз:

@timed
@retry(times=3)
def load(table): ...
 
# то же самое:
load = timed(retry(times=3)(load))

Ближайший к функции декоратор оборачивает её первым, следующий оборачивает уже результат. При вызове load("orders") первым срабатывает код timed, он вызывает обёртку retry, та вызывает исходную функцию. Порядок важен: здесь timed измерит время всех трёх попыток вместе. Если поменять местами, время будет измеряться для каждой попытки отдельно, а retry будет повторять обёртку с замером.

Простой способ запомнить: декораторы это слои луковицы, запись сверху вниз это снаружи внутрь.

Декораторы, которые пригодятся в загрузках

Замер времени, чтобы видеть в логах, какой шаг тормозит:

def timed(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        started = time.perf_counter()
        try:
            return func(*args, **kwargs)
        finally:
            print(f"{func.__name__}: {time.perf_counter() - started:.2f} с")
    return wrapper

finally гарантирует, что время напечатается и при исключении: обычно именно упавший шаг и интересен.

Кэш для справочников, которые читаются из базы по сто раз за прогон: functools.lru_cache или functools.cache из стандартной библиотеки. Это тоже декоратор, и писать свой не нужно. Аргументы функции должны быть хешируемыми, то есть кэшировать lookup(city_id) можно, а lookup(row_dict) нет: привет прошлому уроку.

Декораторы вы будете чаще использовать, чем писать: @task в Airflow, @pytest.fixture, @dataclass, @property, @staticmethod. Все они устроены так же: принимают функцию или класс и возвращают что-то вместо них.

Чего не стоит делать

Декоратор скрывает поведение: читая load("orders"), человек не видит, что там три попытки и замер. Если обёрток становится больше двух, лучше явный код в теле функции или отдельный шаг пайплайна.

Декоратор с состоянием (счётчик вызовов, кэш) делит это состояние между всеми вызовами в процессе. В Airflow воркер живёт долго, и кэш из прошлого запуска задачи может дожить до следующего.

Как отвечать на собеседовании

Что такое декоратор? Функция, которая принимает функцию и возвращает новую, обычно обёртку, которая делает что-то до и после вызова исходной. Запись @dec над def f эквивалентна f = dec(f). Работает благодаря тому, что функции это объекты, а внутренняя функция через замыкание помнит исходную.

Можно ли повесить несколько декораторов и как они применяются? Можно. Применяются снизу вверх (ближайший к функции первым), а при вызове выполняются сверху вниз: @a @b def f это a(b(f)).

Зачем functools.wraps? Чтобы обёртка сохранила имя, docstring и сигнатуру исходной функции. Без него в логах и traceback везде wrapper.

Как сделать декоратор с параметрами? Добавить ещё один уровень: функция с параметрами возвращает декоратор, декоратор возвращает обёртку. @retry(3) это сначала вызов retry(3), а потом применение результата к функции.

Практика

import functools

@count_calls
def parse_amount(raw):
    return int(raw.replace(" ", ""))