Практический пример использования Protocol

Опубликовано 26 February 2024 в Python

Представьте себе ситуацию: у вас есть микросервисы, в каждом своя конфигурация со своим набором параметров; есть библиотечный код, который получает и использует объекты конфигураций. Представили? А теперь в этот код надо завести подсказки типов. И тут представили? Давайте проведем мысленный эксперимент и попробуем это сделать. Вариантов в целом несколько.

  • Базовый класс конфигурации с еще одним ветвлением:
    • затаскивать все возможные параметры;
    • затаскивать только общие для микросервисов параметры.
  • Использовать протоколы для нотации типов.

Затаскивать все возможные варианты параметров в один базовый класс — гарантированный путь к помойке... Затаскивать только общие параметры для микросервисов — получать кучу красного цвета от pyright в общем коде.

Вопрос, зачем нужны конфиги в библиотечном коде, отличный. Правильный ответ: там конфиги не нужны. Там нужны правильные вызовы... Но... Давайте для мысленного эксперимента предположим, что мы знатно косякнули во время написания кода и у нас вместо красивых вызовов в библиотечном коде куча использований объектов конфигураций. Важно: любое совпадение с реальным косяком в реальном коде случайно ;)

Почему Protocol?

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

Протоколы — очень гибкий вариант. Он позволяет добиться еще одной неочевидной выгоды: библиотечный код становится более независимым. Основной недостаток: для большой базы потребует приличного объема кода.

Тут есть противоречие? Я не хочу переписывать код, но мне нужно написать код для внедрения протоколов. На самом деле противоречия нет: - я оставлю код с бизнес-логикой без изменений; - я оставлю интерфейсы классов, методов и функций без изменений; - я напишу в основном дополнительный код для тайп-хинтов.

Возможно, я выкину несколько лишних импортов, что позволит разделить библиотечный код на меньшие независимые узлы.

Что было?

Попробую показать, с какой структурой кода начинаем. Дерево проекта выглядит как-то так.

  • lib
    • base_config.py
    • helper_a.py
    • helper_b.py
  • services
    • service_a
      • config.py
      • app.py
    • service_b
      • config.py
      • app.py

base_config.py содержит в себе класс, описывающий базовые конфиги, которые есть во всех сервисах.

# lib/base_config.py
@dataclass
class BaseConfig:
    db_connection_str: str
    ...

В helper_a.py может лежать код, который требует конфига сервиса.

# lib/helper_a.py
from base_config import BaseConfig

def do_work(config: BaseConfig):
    with db_connection(config.db_connection_str) as connection:
        ...
    ...

Вроде бы все ок. Но вот появилась функция do_more_work, которая используется в сервисах, где есть Redis. Таких сервисов может быть много, но не обязательно все сервисы проекта.

# lib/helper_b.py
from base_config import BaseConfig

def do_more_work(config: BaseConfig):
    redis = get_redis_connection(
        redis_creds=config.redis_creds  # not in BaseConfig
    )
    ...

Придется упаковывать коннект к Redis в общий конфиг для всех сервисов? Не обязательно. Если готовы мириться с тем, что Pyright или другой анализатор подсвечивает вызовы как ошибку, то и переписывать ничего не нужно. Я бы переписал. И вот тут как раз приходят на помощь протоколы.

Используем Protocol

Вместо BaseConfig в каждом случае я делаю класс-наследник от Protocol.

# lib/helper_a.py
from typing import Protocol

class DoWorkConfigProtocol(Protocol):
    db_connection_str: str

def do_work(config: DoWorkConfigProtocol):
    with db_connection(config.db_connection_str) as connection:
        ...
    ...
# lib/helper_b.py
from typing import Protocol

class DoMoreWorkConfigProtocol(Protocol):
    redis_creds: str

def do_more_work(config: DoMoreWorkConfigProtocol):
    redis = get_redis_connection(
        redis_creds=config.redis_creds  # not in BaseConfig
    )
    ...

Я убрал зависимость в хелперах от BaseConfig. Я подсказал анализатору питона, что я жду в config в каждом конкретном случае.

Стал не нужен класс BaseConfig. Смысла в нем больше нет. Где-нибудь в сервисах, когда я буду делать вызовы библиотечных функций, мне анализатор подскажет, все ли в порядке с моим конфигом или мне нужно добавить к нему пропущенный параметр. А главное, он подскажет, что именно мне нужно добавить.

В реальных проектах редко бывает так, что можно ограничиться одним протоколом. Чаще всего нужна какая-то иерархия протоколов. Новые свойства будут добавляться по мере раскручивания стека вызовов или иерархии объектов. Чем глубже по стеку вызов, тем меньше свойств описано в протоколе. Аналогично для классов: в базовом классе минимальный набор, дальше по иерархии протоколы добавляют свойства.

Повторюсь, я бы предпочел явную передачу параметров. Но мы договорились: работаем с тем кодом, что есть. Так что в коде будет что-то вроде такого:

# lib/helper_a.py
from typing import Protocol

class DoWorkConfigProtocol(Protocol):
    db_connection_str: str

def do_work(config: DoWorkConfigProtocol):
    with db_connection(config.db_connection_str) as connection:
        ...
    ...

# lib/helper_b.py
from typing import Protocol
from lib import helper_a

class DoMoreWorkConfigProtocol(
    helper_a.DoWorkConfigProtocol, 
    Protocol
):
    redis_creds: str

def do_more_work(config: DoMoreWorkConfigProtocol):
    redis = get_redis_connection(
        redis_creds=config.redis_creds  # not in BaseConfig
    )
    do_work(config)
    ...

Заключение

Работа с протоколами довольно запутанная. Есть много моментов, на которые стоит обращать внимание. Не могу сказать, что я добился полного их понимания. Мне помогает опыт работы с Java. Протокол — практически один в один интерфейс в Java. И используется преимущественно для тех же целей.

Не могу оставить вас без полезных ссылок для дальнейшего изучения:

  • Лучшее описание протоколов, что я встречал, — это статья Protocols and structural subtyping из документации к mypy.
  • Не обойтись без PEP 544 — описание протоколов, как они приняты и реализованы в Python.

А вы используете в своем коде протоколы? Как? Было бы интересно услышать ваше мнение.

---
Возник вопрос? Мне всегда можно написать в Twitter: avkorablev