Декораторы: добавляем поведение функции, не трогая её

Что мы изучим?#

• Что такое декоратор (Decorator) и чем он полезен.
• Как построить декоратор: функция, которая получает функцию и возвращает обёртку (wrapper).
• Синтаксис @.
• Обёртку, подходящую для любой функции, с *args и **kwargs.

Зачем это нужно?#

Мы хотим печатать сообщение при каждом вызове функции — для отладки:

def add(a, b):
    print("Calling add")
    return a + b

def multiply(a, b):
    print("Calling multiply")
    return a * b

Проблемы:

• Одна и та же строка копируется в каждую функцию — снова дублирование, которое запрещает DRY.
• Логика смешивается со вспомогательным кодом, который к ней не относится.
• Чтобы убрать печать, нужно пройти по всем функциям и удалить её.

Декоратор позволяет написать дополнительное поведение один раз и надеть его на любую функцию одной строкой.

Объяснение с основ#

Из чего построен декоратор#

Декоратор сочетает две вещи, которые мы изучили в предыдущих уроках:

• Функцию высшего порядка — он получает функцию и возвращает функцию.
• Замыкание — функция, которую он возвращает, помнит исходную функцию.

Структура всегда одна и та же:

1. Декоратор получает исходную функцию (func).
2. Внутри него определяется новая функция — wrapper (обёртка), — которая делает что-то дополнительное и вызывает func.
3. Декоратор возвращает wrapper — без скобок.

Синтаксис @#

@log_call
def add(a, b):
    return a + b

Это абсолютно то же самое, что написать после определения:

add = log_call(add)

То есть имя add теперь указывает на обёртку, а обёртка вызывает исходную add.

Синтаксис#

def my_decorator(func):              # 1. receive the original function
    def wrapper(*args, **kwargs):    # 2. a new function that wraps it
        # ... something before ...
        result = func(*args, **kwargs)
        # ... something after ...
        return result                # return the original result
    return wrapper                   # 3. return the wrapper (no parentheses)

@my_decorator                        # same as: my_function = my_decorator(my_function)
def my_function():
    ...

Первый пример кода#

week-2/log_call.py:

def log_call(func):
    def wrapper(a, b):
        print("Calling", func.__name__)
        return func(a, b)
    return wrapper

@log_call
def add(a, b):
    return a + b

@log_call
def multiply(a, b):
    return a * b

print(add(2, 3))
print(multiply(4, 5))

Вывод:

Calling add
5
Calling multiply
20

Разбор кода построчно#

Ещё один пример#

Обёртка, подходящая для любой функции. Предыдущая обёртка работает только с функциями, которые получают ровно a, b. С *args и **kwargs она принимает любые параметры и передаёт их дальше. week-2/log_any.py:

def log_call(func):
    def wrapper(*args, **kwargs):
        print("Calling", func.__name__, "with", args, kwargs)
        return func(*args, **kwargs)
    return wrapper

@log_call
def greet(name, greeting="Hello"):
    return greeting + ", " + name

@log_call
def total(*prices):
    return sum(prices)

print(greet("Dana", greeting="Hi"))
print(total(10, 20, 30))

Вывод:

Calling greet with ('Dana',) {'greeting': 'Hi'}
Hi, Dana
Calling total with (10, 20, 30) {}
60

Полезный пример: измерение времени. week-2/timed.py:

import time

def timed(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        seconds = round(time.time() - start, 2)
        print(func.__name__, "took", seconds, "seconds")
        return result
    return wrapper

@timed
def slow_sum(n):
    time.sleep(0.5)
    return sum(range(n))

print(slow_sum(1000))

Вывод:

slow_sum took 0.5 seconds
499500

(Точное время немного меняется от запуска к запуску.) Тот же @timed можно надеть на любую функцию программы — не меняя в ней ни одной строки.

Типичная ошибка#

Две ошибки, которые случаются почти у всех, кто пишет первый декоратор:

def log_call(func):
    def wrapper(*args, **kwargs):
        print("Calling", func.__name__)
        func(*args, **kwargs)
    return wrapper()

1. Нет return внутри обёртки. Исходная функция выполняется, но её результат теряется — и декорированная функция всегда возвращает None. Исправление: return func(*args, **kwargs).
2. return wrapper() со скобками. Это вызывает обёртку уже при определении и возвращает её результат вместо функции. Исправление: return wrapper.

Правило: обёртка возвращает результат, а декоратор возвращает обёртку.

Практика#

Упражнения опираются только на то, что пройдено в этом уроке и в предыдущих. Сначала попробуйте сами, и только потом открывайте решение.

Упражнение 1: Первый декоратор#

Напишите декоратор с именем log_call, который печатает Calling <name> перед каждым вызовом функции, и примените его через @ к add и к multiply.

print(add(2, 3))
print(multiply(4, 5))

Ожидаемый вывод:

Calling add
5
Calling multiply
20

def log_call(func):
    def wrapper(a, b):
        print("Calling", func.__name__)
        return func(a, b)
    return wrapper


@log_call
def add(a, b):
    return a + b


@log_call
def multiply(a, b):
    return a * b


print(add(2, 3))
print(multiply(4, 5))

Calling add
5
Calling multiply
20

Упражнение 2: Обёртка, подходящая для любой функции#

Предыдущая обёртка работает только с функциями, которые получают ровно два параметра. Перепишите log_call так, чтобы он работал с любой функцией.

@log_call
def greet(name, greeting="Hello"):
    return greeting + ", " + name

@log_call
def total(*prices):
    return sum(prices)

print(greet("Dana", greeting="Hi"))
print(total(10, 20, 30))

Ожидаемый вывод:

Calling greet
Hi, Dana
Calling total
60

def log_call(func):
    def wrapper(*args, **kwargs):
        print("Calling", func.__name__)
        return func(*args, **kwargs)
    return wrapper


@log_call
def greet(name, greeting="Hello"):
    return greeting + ", " + name


@log_call
def total(*prices):
    return sum(prices)


print(greet("Dana", greeting="Hi"))
print(total(10, 20, 30))

Calling greet
Hi, Dana
Calling total
60

Упражнение 3: Проверка ввода#

Напишите декоратор с именем positive_only для функций, которые получают одно число: если число отрицательное, он выбрасывает ValueError и вообще не вызывает функцию. Сообщение включает имя функции (func.__name__) и полученное число, как в ожидаемом выводе.

@positive_only
def square_root(n):
    return n ** 0.5

print(square_root(16))
try:
    square_root(-4)
except ValueError as error:
    print("error:", error)

Ожидаемый вывод:

4.0
error: square_root needs a positive number, got -4

def positive_only(func):
    def wrapper(n):
        if n < 0:
            raise ValueError(func.__name__ + " needs a positive number, got " + str(n))
        return func(n)
    return wrapper


@positive_only
def square_root(n):
    return n ** 0.5


print(square_root(16))
try:
    square_root(-4)
except ValueError as error:
    print("error:", error)

4.0
error: square_root needs a positive number, got -4

Упражнение 4: Подсчёт вызовов#

Напишите декоратор с именем count_calls, который считает, сколько раз была вызвана функция. При каждом вызове, до выполнения функции, он печатает имя функции и номер вызова — как в ожидаемом выводе.

@count_calls
def say_hi():
    print("Hi")

say_hi()
say_hi()

Ожидаемый вывод:

say_hi call #1
Hi
say_hi call #2
Hi

def count_calls(func):
    count = 0

    def wrapper(*args, **kwargs):
        nonlocal count
        count += 1
        print(func.__name__, "call #" + str(count))
        return func(*args, **kwargs)

    return wrapper


@count_calls
def say_hi():
    print("Hi")


say_hi()
say_hi()

say_hi call #1
Hi
say_hi call #2
Hi

Упражнение 5: Две ошибки#

def shout(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        result.upper()
    return wrapper()

Декоратор должен возвращать результат функции заглавными буквами. Найдите две ошибки и исправьте их.

@shout
def greet(name):
    return "hello " + name

print(greet("dana"))

Ожидаемый вывод:

HELLO DANA

def shout(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        return result.upper()
    return wrapper


@shout
def greet(name):
    return "hello " + name


print(greet("dana"))

HELLO DANA

Итоги#

• Декоратор = функция, которая получает функцию, заворачивает её в обёртку с дополнительным поведением и возвращает обёртку.
• Он строится из функции высшего порядка и замыкания — двух тем предыдущих уроков.
• @decorator над def = func = decorator(func).
• С *args и **kwargs обёртка подходит для любой функции.
• Обёртка возвращает результат (return func(...)), а декоратор возвращает обёртку (return wrapper, без скобок).

Официальная документация#

• Python Glossary — decorator
• Function definitions
• functools.wraps