Как создать собственные Python-декораторы и правильно их использовать

Статья рассчитана на тех, кто владеет основами Python, знаком с декораторами и хочет научиться создавать собственные декораторы для повышения качества кода. Если вы забыли, что такое декораторы, — повторите тему по первым разделам статьи.

Михаил Залешин

К.ф.-м.н., Data Scientist, наставник и эксперт в Нетологии по курсам «Программист на JavaScript» и «Data Science: Рекомендательные системы»

Что такое декораторы в Python

Декоратор — конструкция языка Python для расширения возможностей функций и классов без изменения их кода.

Есть смартфон. Сделаем его устойчивым к падениям — наденем на него чехол. Чехол не изменяет прежние возможности смартфона, и добавляет к нему качество — ударопрочность. Чехол — декоратор смартфона.

Один и тот же чехол подходит для всех смартфонов нужной модели. Универсальность — важное свойство декораторов.

Анатомия декоратора в Python

Создадим декоратор @hello_decorator:

			from functools import wraps

def hello_decorator(f):
   @wraps(f)
   def wrapper(*args, **kwargs):
       print('Hello from decorator!')
       return f(*args, **kwargs)

   return wrapper
		

Декоратор в Python — функция, которая принимает функцию/класс и возвращает функцию/класс. В примере выше декоратор hello_decorator() принимает функцию f(), и возвращает функцию wrapper().

Пояснения:

1. При объявлении функции wrapper() используем встроенный в Python декоратор @wraps. Этот декоратор копирует свойства __name__, __doc__ и другие из функции f() в функцию wrapper(), чтобы при отладке программы все выглядело так, будто wrapper() и есть функция f();
2. Аргументы функции wrapper(): *args и **kwargs. Аргумент *args собирает позиционные аргументы, а **kwargs — именованные. Например, в вызове: wrapper(1, ‘a’, x=5, y=None) значение args — кортеж (1, ‘a’), а kwargs — словарь {‘x’: 5, ‘y’: None}. Если позиционных аргументов при вызове функции нет, args — пустой, и если нет именованных аргументов, пустой — kwargs;
3. В первой строке тела функции wrapper() в консоль выводится «Hello from decorator!» — единственный «побочный эффект» декоратора. Далее вызывается декорируемая функция f();
4. Функция f() внутри wrapper() принимает параметры *args и **kwargs. Операторы * и ** перед именами параметров в вызове функции имеют противоположный эффект случаю, когда их используют при объявлении аргументов. Например, если args=(1, ‘a’) и kwargs={‘x’: 5, ‘y’: None}, вызов f(*args, **kwargs) равносилен: f(1, ‘a’, x=5, y=None). Комбинация «звездочных» аргументов и «звездочных» операторов позволяет универсально передать аргументы из функции wrapper() в f();
5. В последней строке функции hello_decorator() возвращаем функцию-обертку wrapper(). Так мы указываем, что нужно подставить на место декорируемой функции f(). Вызывать функцию wrapper() не нужно — возвращаем саму функцию.

Применим декоратор @hello_decorator к функции sum2():

			@hello_decorator
def sum2(a, b):
   return a + b
		

Функция sum2() принимает два аргумента и возвращает их сумму. Декорированный sum2() дополнительно выводит на консоль «Hello from decorator!».

Синтаксис @hello_decorator введен в Python для удобства, и равносилен такой записи:

			def sum2(a, b):
   return a + b

sum2 = hello_decorator(sum2)
		

Настраиваемый логгер-декоратор

Наметим логирующий декоратор @Logger для функций с учетом следующих пожеланий:

1. Логи отправляются на консоль;
2. Вложенность: если логируемая функция вызывает другую логируемую функцию, делаем при выводе на консоль отступ для последней;
3. Уровни логов DEBUG, INFO, CRIT: при декорировании можно указать уровень лога функции, который работает в комбинации с настройкой детальности отображения логов. Если уровень лога функции выше или равен текущей детальности, отображаем лог, иначе — игнорируем;
4. Для логирования вызова декорированной функции используем шаблон: LOG_LEVEL [TIMESTAMP] FUNC_NAME(ARGS, KWARGS);
5. Логирование исключений: если функция выбрасывает исключение, логируем его с уровнем CRIT;
6. Внутренние логи: логгер должен работать и в режиме декоратора, и как функция для вывода в лог произвольных сообщений.

Начнем с примера использования. Так мы не перегружаем внимание внутренней сложностью и повышаем шансы создать удачный интерфейс модуля. На этом принципе основана разработка через тестирование — test-driven development (TTD).

Что хотим получить:

			logger = Logger(verbosity=Logger.LogLevel.DEBUG)

@logger()
def load_data(url):
   from string import ascii_lowercase
   import random
   return random.choices(ascii_lowercase, k=3)

@logger(log_level=Logger.LogLevel.DEBUG)
def check_value(value):
   logger.log_msg(
       Logger.LogLevel.DEBUG,
       'Doing some important stuff...'
   )
   return True

@logger()
def process_value(value):
   if check_value(value):
       return value.upper()
   return value

@logger()
def save_data(url, data):
   raise Exception('Could not save data :(')

@logger()
def main():
   data = load_data('example.com')
   data = [*map(process_value, data)]
   save_data('example.com', data)
		

После вызова main() хотим увидеть в консоли:

			⠀INFO [2022-09-29 17:14:26] main((), {})
 INFO [2022-09-29 17:14:26] -> load_data(('example.com',), {})
 INFO [2022-09-29 17:14:26] -> process_value(('z',), {})
DEBUG [2022-09-29 17:14:26] ---> check_value(('z',), {})
DEBUG [2022-09-29 17:14:26] -----> Doing some important stuff...
 INFO [2022-09-29 17:14:26] -> process_value(('u',), {})
DEBUG [2022-09-29 17:14:26] ---> check_value(('u',), {})
DEBUG [2022-09-29 17:14:26] -----> Doing some important stuff...
 INFO [2022-09-29 17:14:26] -> process_value(('f',), {})
DEBUG [2022-09-29 17:14:26] ---> check_value(('f',), {})
DEBUG [2022-09-29 17:14:26] -----> Doing some important stuff...
 INFO [2022-09-29 17:14:26] -> save_data(('example.com', ['Z',
                                          'U', 'F']), {})
 CRIT [2022-09-29 17:14:26] ---> Could not save data :(
		

Наблюдения по поводу декоратора @Logger:

1. Logger — не функция, а класс. Экземпляры этого класса можно вызвать. Результат вызова — декоратор;
2. У класса Logger есть вложенный класс-перечисление LogLevel с полями DEBUG, INFO, CRIT;
3. Детальность логирования verbosity определяется в конструкторе класса Logger, уровень логирования функции log_level задается при ее декорировании;

У класса Logger есть метод log_msg(), который можно использовать напрямую внутри функций.

Напишем скелет класса Logger:

			from contextlib import contextmanager
from functools import wraps
from datetime import datetime
from enum import Enum

class Logger:
   class LogLevel(Enum):
       DEBUG = 0
       INFO = DEBUG + 1
       CRIT = INFO + 1

       # Помогает принять решение о пропуске лога,
       # который имеет уровень ниже, чем установленная
       # детальность логирования
       def should_skip(self, verbosity):
           return self.value < verbosity.value

   # Определяет ширину отступа слева для каждого уровня лога
   PREFIX_OFFSET = 2

   def __init__(self, verbosity=LogLevel.INFO):
       pass

   # Этот метод делает экземпляры класса вызываемыми
   # Используем его для декорирования
   # Параметр should_rethrow_exceptions определяет, нужно
   # ли подавлять исключения или передавать их на уровень выше
   def __call__(self, log_level=LogLevel.INFO,
                should_rethrow_exceptions=False):
       pass

   def log_msg(self, log_level, msg):
       pass

   def log_func(self, log_level, f, args, kwargs):
       pass

   # Контекстный менеджер для контроля глубины лога
   @contextmanager
   def _deeper(self):
       self._depth += 1
       try:
           yield
       finally:
           self._depth -= 1
		

Конструктор класса Logger:

			def __init__(self, verbosity=LogLevel.INFO):
   self._depth = 0 # для отслеживания глубины
   self.verbosity = verbosity
		

Метод Logger.log_msg():

			def log_msg(self, log_level, msg):
   # Пропускаем логи с уровнем ниже уровня детализации
   if log_level.should_skip(self.verbosity):
       return

   timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
   prefix = f'{log_level.name:>5} [{timestamp}] '
   prefix += '-' * (Logger.PREFIX_OFFSET * self._depth - 1)
   if self._depth > 0:
       prefix += '>'

   print(f'{prefix} {msg}')
		

Вспомогательный метод Logger.log_func():

			def log_func(self, log_level, f, args, kwargs):
   self.log_msg(log_level, f'{f.__name__}({args}, {kwargs})')
		

Ключевой метод-декоратор Logger.__call__():

			def __call__(self, log_level=LogLevel.INFO,
             should_rethrow_exceptions=False):
   # Декоратор должен принимать один аргумент — декорируемую
   # функцию
   # Если декоратору нужны параметры, используем замыкание с
   # помощью внешней функции
   def decorator(f):
       @wraps(f)
       def wrapper(*args, **kwargs):
           self.log_func(log_level, f, args, kwargs)

           try:
               # Используем менеджер контекста
               with self._deeper():
                   return f(*args, **kwargs)

           except Exception as e:
               # Логируем исключения
               if log_level.should_skip(self.verbosity):
                   self.log_func(Logger.LogLevel.CRIT, f, args,
                                 kwargs)

               with self._deeper():
                   self.log_msg(Logger.LogLevel.CRIT, str(e))

               if should_rethrow_exceptions:
                   raise e

       return wrapper

   return decorator
		

Выводы

1. Проектировать проще сверху-вниз: сначала решите, как хотите пользоваться модулем, а потом реализуйте его;
2. Если декоратор требует централизованного контроля, используйте класс;
3. Если декоратору нужны дополнительные параметры, используйте замыкание — оберните функцию-декоратор во внешнюю функцию (см. Logger.__call__());
4. Экземпляры классов можно использовать как функции за счет __call__();
5. Менеджеры контекста удобны, когда нужно «подчищать» состояние.