Генерация документации из моделей с помощью Pydantic
Привет, меня зовут Роман Брылунов, я QA Automation в команде сервиса транспорта 2ГИС. Мы автоматизируем сервисы транспортных сценариев, таких как построение маршрутов для разных видов транспорта, построение пешеходных маршрутов, решение задачи коммивояжера. Основная часть наших тестов — функциональные тесты логики приложения.
С сервисами общаемся по HTTP, но есть и несколько внутренних сервисов со взаимодействием по gRPC. Все ответы, полученные от сервисов, мы предварительно валидируем перед обработкой в тесте. Для валидации используем библиотеку Pydantic. Это позволяет нам описывать формат взаимодействия с помощью моделей и обрабатывать ответы в виде Python-объектов вместо словаря после стандартного парсинга JSON. Тесты встроены в CI, успешное прохождение тестов является блокирующим условием для влития кода. Таким образом, чтобы внести изменения в API, необходимо актуализировать модели в тестах. В противном случае ответ от сервиса не пройдёт валидацию при прогоне тестов.
Мы стараемся всесторонне подходить к контролю качества продукта, в том числе хотим, чтобы у продукта была красивая и актуальная документация. У наших сервисов была документация, но она обновлялась вручную. А у нас есть модели запросов и ответов, которые точно соответствуют текущему формату API. Мы подумали, что будет здорово использовать эти модели для формирования документации. Попробовали и успешно внедрили. Расскажу о том, как можно это сделать.
Базовая реализация
Представим, что у нас есть некий сервис, реализующий REST API по протоколу HTTP. Для проверки его работы у нас есть автотесты, которые используют модели для валидации и обработки ответов от сервиса. На тестах фокусироваться не будем, просто считаем, что они есть и их запуск встроен в CI. Обратим внимание именно на модели.
Для начала создадим базовый класс для всех наших моделей, который будет отвечать за настройку поведения. Будем использовать строгие модели, не допускающие наличие неописанных полей.
from pydantic import BaseModel
from pydantic import ConfigDict
class NavigationApiBaseModel(BaseModel):
model_config = ConfigDict(
use_enum_values=True, # используем значение из Enum
validate_assignment=True, # включаем валидацию при изменении
extra='forbid', # запрещаем не описанные в модели поля
)Для примера рассмотрим некий эндпоинт, который принимает список точек и отдаёт список найденных маршрутов. Напишем модели для запроса и ответа, унаследовав базовый класс в качестве родительского.
from enum import Enum
from base_models import NavigationApiBaseModel
class RequestPoint(NavigationApiBaseModel):
lat: float
lon: float
class NavigationRequest(NavigationApiBaseModel):
points: list[RequestPoint]
class NavigationStatus(Enum):
SUCCESS = 'success'
FAIL = 'fail'
class NavigationResponse(NavigationApiBaseModel):
status: NavigationStatus
duration: int
length: intДля тестов этого будет достаточно, но мы ведь хотим использовать модели и для генерации документации. Один из удобных вариантов хранения документации — OpenAPI-схема. Из коробки Pydantic умеет генерировать только JSON-схему для модели, а OpenAPI-схема содержит в себе не только описание самих моделей, но еще и описание эндпоинтов. Благо, мы не единственные, кто задумался над генерацией OpenAPI-схемы из Pydantic-моделей — это умеет библиотека openapi-pydantic.
Опишем наш эндпоинт для построения маршрутов в соответствии со спецификацией OpenAPI в виде словаря в python:
info = {
'openapi': '3.1.0',
'info': {'title': 'Navigation API', 'version': '6.0.0'},
'servers': [{'url': 'Не заполнены только схема запроса и ответа. Добавим ссылку на наши модели вместо схем:
# ******************** Новый фрагмент ********************
# импортируем класс для добавления ссылки на модель
from openapi_pydantic.util import PydanticSchema
# и наши модели
from models import NavigationRequest
from models import NavigationResponse
# ******************** Конец нового фрагмента ********************
info = {
'openapi': '3.1.0',
'info': {'title': 'Navigation API', 'version': '6.0.0'},
'servers': [{'url': 'Добавим ещё немного кода для генерации схемы:
# ******************** Новый фрагмент ********************
# еще пара импортов из библиотеки для генерации схемы
from openapi_pydantic import OpenAPI
from openapi_pydantic.util import construct_open_api_with_schema_class
# ******************** Конец нового фрагмента ********************
from openapi_pydantic.util import PydanticSchema
from models import NavigationRequest
from models import NavigationResponse
info = {
'openapi': '3.1.0',
'info': {'title': 'Navigation API', 'version': '6.0.0'},
'servers': [{'url': 'https://example.com/'}],
'paths': {
'/navigation': {
'post': {
'summary': 'Построение маршрута',
'description': 'API для прокладывания маршрутов.',
'requestBody': {
'description': 'Параметры запроса',
'content': {
'*/*': {
'schema': PydanticSchema(schema_class=NavigationRequest)
}
},
'required': True,
},
'responses': {
'200': {
'description': 'Набор маршрутов',
'content': {
'application/json': {
'schema': PydanticSchema(schema_class=NavigationResponse)
}
},
},
},
}
},
}
}
# ******************** Новый фрагмент ********************
open_api_without_components = OpenAPI(**info)
open_api = construct_open_api_with_schema_class(open_api_without_components)
with open('navigation_api.json', 'w', encoding='utf-8') as out_file:
out_file.write(open_api.model_dump_json(by_alias=True, exclude_none=True, indent=2))
# ******************** Конец нового фрагмента ********************Запускаем файл и получаем в файле navigation_api.json готовую OpenAPI-схему:
{
"openapi": "3.1.0",
"info": {
"title": "Navigation API",
"version": "6.0.0"
},
"servers": [
{
"url": "Схема есть, но она пока не подходит для использования в качестве документации. По крайней мере в том виде, в котором мы бы хотели видеть документацию — не хватает описаний.
Добавляем описания
Для простых полей описание можно добавить через стандартный для Pydantic-класс Field.
from pydantic import Field
from base_models import NavigationApiBaseModel
class RequestPoint(NavigationApiBaseModel):
lat: float = Field(description='Градусы долготы.')
lon: float = Field(description='Градусы широты.')Для каждого класса с моделью тоже добавим описание через docstring для класса.
from enum import Enum
from pydantic import Field
from base_models import NavigationApiBaseModel
class RequestPoint(NavigationApiBaseModel):
"""Координаты точки маршрута."""
lat: float = Field(description='Градусы долготы.')
lon: float = Field(description='Градусы широты.')
class NavigationRequest(NavigationApiBaseModel):
"""Запрос для построения маршрута."""
points: list[RequestPoint]
class NavigationStatus(Enum):
"""Статус маршрута."""
SUCCESS = 'success'
FAIL = 'fail'
class NavigationResponse(NavigationApiBaseModel):
"""Результат построения маршрута."""
status: NavigationStatus
duration: int = Field(description='Время в пути в секундах.')
length: int = Field(description='Длина маршрута в метрах.')Таким образом поля, ссылающиеся на другую модель, тоже получают описания. Вот так это будет выглядеть в схеме:
"components": {
"schemas": {
"NavigationRequest": {
"properties": {
"points": {
"items": {
"$ref": "#/components/schemas/RequestPoint"
},
"type": "array",
"title": "Points"
}
},
"additionalProperties": false,
"type": "object",
"required": [
"points"
],
"title": "NavigationRequest",
"description": "Запрос для построения маршрута."
},
"NavigationResponse": {
"properties": {
"status": {
"$ref": "#/components/schemas/NavigationStatus"
},
"duration": {
"type": "integer",
"title": "Duration",
"description": "Время в пути в секундах."
},
"length": {
"type": "integer",
"title": "Length",
"description": "Длина маршрута в метрах."
}
},
"additionalProperties": false,
"type": "object",
"required": [
"status",
"duration",
"length"
],
"title": "NavigationResponse",
"description": "Результат построения маршрута."
},
"NavigationStatus": {
"type": "string",
"enum": [
"success",
"fail"
],
"title": "NavigationStatus",
"description": "Статус маршрута."
},
"RequestPoint": {
"properties": {
"lat": {
"type": "number",
"title": "Lat",
"description": "Градусы долготы."
},
"lon": {
"type": "number",
"title": "Lon",
"description": "Градусы широты."
}
},
"additionalProperties": false,
"type": "object",
"required": [
"lat",
"lon"
],
"title": "RequestPoint",
"description": "Координаты точки маршрута."
}
}
}Схема с описаниями выглядит уже значительно более информативно. Дальше можно заняться кастомизацией схемы под конкретные задачи.
Кастомизируем схему
Допустим, если захочется опубликовать схему в качестве публичной документации, то чтобы не светить внутренние поля, надо научиться помечать их как приватные. Стандартный Field подходящего атрибута не имеет, но позволяет передавать дополнительную информацию через атрибут json_schema_extra.
class NavigationResponse(NavigationApiBaseModel):
"""Результат построения маршрута."""
status: NavigationStatus
duration: int = Field(description='Время в пути в секундах.')
length: int = Field(description='Длина маршрута в метрах.')
calculate_duration: int | None = Field(
None,
description='Время вычисления маршрута в секунда.',
json_schema_extra={'is_private': True}
)Чтобы фильтровать приватные поля при экспорте схемы, у модели надо определить метод __get_pydantic_json_schema__ (часть старого кода скрыта за …)
from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import GetJsonSchemaHandler
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema as cs
class NavigationApiBaseModel(BaseModel):
...
@classmethod
def __get_pydantic_json_schema__(cls, core_schema: cs.CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
json_schema = handler(core_schema)
json_schema = handler.resolve_ref_schema(json_schema)
# Обходим поля объекта
private_fields = []
if 'properties' in json_schema:
# Собираем приватные поля
for key, prop in json_schema['properties'].items():
if prop.get('is_private'):
private_fields.append(key)
# И исключаем их
for private_key in private_fields:
json_schema['properties'].pop(private_key)
return json_schemaГенерируем схему и убеждаемся, что приватного поля в ней нет. По аналогии через метод __get_pydantic_json_schema__ можно добавить, например, возможность генерировать схему для разных языков, указывая для каждого языка своё описание. Или убирать из схемы ненужные атрибуты.
Ещё важно генерировать детальные описания допустимых значений для Enum, чтобы красиво отобразить их в публичной документации. Для этого сделаем кастомный Enum, который позволит хранить описание для значения. А для передачи таких значений с описанием заведём отдельный класс.
from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
@dataclass()
class DescribedEnumValue:
"""Класс для создания значений у Enum с описанием."""
value: any
description: str = ''
class DescribedEnum(Enum):
"""Enum с возможностью указать описание."""
def __new__(cls, value: any | DescribedEnumValue):
obj = object.__new__(cls)
# если получили в качестве значения объект с описанием – сохраняем
# описание и извлекаем действительное значение для элемента Enum
if isinstance(value, DescribedEnumValue):
obj._value_ = value.value
obj._description = value.description
else:
obj._value_ = value
obj._description = ''
return objПри генерации схемы научим класс собирать эти описания в description самого Enum (часть старого кода скрыта за …).
from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
...
class DescribedEnum(Enum):
"""Enum с возможностью указать описание."""
...
@classmethod
def __get_pydantic_json_schema__(cls, core_schema: cs.CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
"""Модификация схемы при экспорте."""
json_schema = handler(core_schema)
json_schema = handler.resolve_ref_schema(json_schema)
# Получаем красивое описание с HTML-разметкой списка значений
members_description = ''
for member in cls:
if member._description:
members_description += f'{member._value_} - {member._description}, '
# Добавляем к описанию информацию о значениях
if members_description:
json_schema['description'] = f'{json_schema["description"]}: {members_description}
'
return json_schemaДобавляем описания для нашего Enum
from base_models import DescribedEnumValue
from base_models import DescribedEnum
class NavigationStatus(DescribedEnum):
"""Статус маршрута"""
SUCCESS = DescribedEnumValue('success', 'Успешное построение')
FAIL = DescribedEnumValue('fail', 'Ошибка при построении')В итоге получаем схему:
"NavigationStatus": {
"type": "string",
"enum": [
"success",
"fail"
],
"title": "NavigationStatus",
"description": "Статус маршрута: - success - Успешное построение,
- fail - Ошибка при построении,
"
},В Swagger схема выглядит так:
Пример Enum с описанием
Больше не надо гадать по значению, что конкретно оно означает. Все подробно расписано в описании.
Результат
На сегодняшний день мы уже используем генерируемые из моделей OpenAPI-схемы для публичной документации навигационных сервисов. Посмотреть можно, например, здесь. Есть некоторые шероховатости с обработкой опциональных полей, их хотим поправить в ближайшее время. Если тоже решитесь использовать модели для генерации документации, вот какие преимущества можно получить:
OpenAPI-схема всегда точно соответствует текущему состоянию кода.
Актуальность схемы удобно и легко поддерживать. Редактировать модели удобнее, чем OpenAPI-схему в виде JSON.
С небольшими доработками можно добавить поддержку разных языков. При этом не придётся иметь две/три/… схем с переводами на разные языки и поддерживать их консистентность.
Можно шарить модели и Enum-ы между схемами, что тоже облегчает поддержку.
Если остались вопросы, обязательно пишите — отвечу.
Больше про RnD — в нашем Телеграм-канале, а ещё у QA есть отдельный чат. Захочешь работать в 2ГИС, загляни в вакансии, у нас хорошо!
