Как делать разную валидацию в зависимости от значения одного из полей

Вопрос следующий:

У нас есть JSON со следующей структурой:

{
    key1: [
        {
         type_x: str,
         atr1: str,
         atr2: int
        },
        {
         type_x: str,
         atr1: str
        },
        {
         type_x: str,
         atr1: str,
         atr2: int,
         atr3: int
        },
    ],
    key2: [
        {
         atr1: str,
         atr2: int
        }
    ]
}

Мне нужно создать модель для валидации, но для различных элементов списка словарей по ключу key1 надо применить различную валидацию в зависимости от значения ключа type_x. Как видите количество параметров в элементах списка различное и скажем так требуемые параметры зависят от значения параметра type_x. Например если значение type_x = 'base', то для него модель:

class MyBase(BaseClass):
    type_x: str
    atr1: str
    atr2: int

если значение type_x = 'root':

class MyRoot(BaseClass):
    type_x: str
    atr1: str
    atr3: int

Как прописать одинаковую валидацию для каждого элемента списка key1 я разобрался:

class MyClass(BaseClass):
    type_x: str
    atr1: str
    atr2: int

А вот как сделать различные требуемые параметры в зависимости от значения type_x? И конечно в идеале это сделать в рамках одной общей модели, чтобы не плодить отдельные модели для каждого варианта значения type_x. Спасибо.


Ответы (2 шт):

Автор решения: insolor

Для проверки по нескольким полям используйте root_validator:

from pydantic import BaseModel, root_validator


class MyClass(BaseModel):
    type_x: str
    atr1: str
    atr2: int
    
    @root_validator
    def check_depending_on_type_x(cls, values):
        # values - словарь значений атрибутов,
        # значения получаем по строковому имени атрибутов
        type_x = values.get("type_x")
        
        # Тут пишите ваше условие, в зависимости от которого должны выполняться разные проверки
        if type_x == "something":  
            if values.get("atr1") != "qwerty":  # Меняете на нужную проверку
                raise ValueError("Ошибка 1")
        else:
            if values.get("atr2") != 42:  # Меняете на вашу проверку (другую)
                raise ValueError("Ошибка 2")
        
        # Если все хорошо, возвращаем values
        return values

Обновление. Если поле может отсутствовать в JSON, "оберните" его тип в Optional (т.е. буквально "не обязательное"), отсутствующие поля будут инициализироваться значением None:

from pydantic import BaseModel
from typing import Optional


class MyRoot(BaseModel):
    type_x: str
    atr1: Optional[str]
    atr2: Optional[int]
    atr3: Optional[int]


print(repr(
    MyRoot.parse_obj({"type_x": "123", "atr1": "asdf", "atr2": 1})
))

print(repr(
    MyRoot.parse_obj({"type_x": "345", "atr1": "qwer"})
))

print(repr(
    MyRoot.parse_obj({"type_x": "789", "atr1": "qwerty", "atr2": 2, "atr3": 3})
))

Вывод:

MyRoot(type_x='123', atr1='asdf', atr2=1, atr3=None)
MyRoot(type_x='345', atr1='qwer', atr2=None, atr3=None)
MyRoot(type_x='789', atr1='qwerty', atr2=2, atr3=3)

Поверх этого уже можно прикрутить валидацию заполненности полей в зависимости от значения type_x аналогично тому, как показано в первой части ответа (если поле должно быть заполнено, проверяйте, что оно is not None).

Если потом нужна будет конвертация в JSON, и нужно чтобы "пустые" поля в JSON не попадали, конвертируйте в json с параметром exclude_unset=True:

my_root = MyRoot.parse_obj({"type_x": "123", "atr1": "asdf", "atr2": 1})
print(repr(my_root))  # MyRoot(type_x='123', atr1='asdf', atr2=1, atr3=None)
print(my_root.json())  # {"type_x": "123", "atr1": "asdf", "atr2": 1, "atr3": null}
print(my_root.json(exclude_unset=True))  # {"type_x": "123", "atr1": "asdf", "atr2": 1}
→ Ссылка
Автор решения: insolor

Реализация через отдельные классы, конкретный тип определяется по значению одного из полей. В документации по Pydantic это называется Discriminated Unions (a.k.a. Tagged Unions).

from pydantic import BaseModel, ValidationError, Field, parse_obj_as
from typing import Literal, Union
from typing_extensions import Annotated


class MyBase(BaseModel):
    type_x: Literal["base"]
    atr1: str
    atr2: int


class MyRoot(BaseModel):
    type_x: Literal["root"]
    atr1: str
    atr3: int


CommonModel = Annotated[Union[MyBase, MyRoot], Field(discriminator='type_x')]


examples = [
    {"type_x": "base", "atr1": "asdf", "atr2": 1},  # Корректная модель base
    {"type_x": "root", "atr1": "asdf", "atr3": 1},  # Корректная модель root
    {"type_x": "foobar"},  # Неизвестный тип
    {"type_x": "base", "atr1": "qwer"},  # Отсутствует одно поле для типа base
    {"type_x": "root"},  # Отсутствуют два поле для типа root
]


for example in examples:
    print(example)
    try:
        print(repr(
            parse_obj_as(CommonModel, example)
        ))
    except ValidationError as ex:
        print(ex)
    print()

Вывод:

{'type_x': 'base', 'atr1': 'asdf', 'atr2': 1}
MyBase(type_x='base', atr1='asdf', atr2=1)

{'type_x': 'root', 'atr1': 'asdf', 'atr3': 1}
MyRoot(type_x='root', atr1='asdf', atr3=1)

{'type_x': 'foobar'}
1 validation error for ParsingModel[typing_extensions.Annotated[Union[__main__.MyBase, __main__.MyRoot], FieldInfo(default=PydanticUndefined, discriminator='type_x', extra={})]]
__root__
  No match for discriminator 'type_x' and value 'foobar' (allowed values: 'base', 'root') (type=value_error.discriminated_union.invalid_discriminator; discriminator_key=type_x; discriminator_value=foobar; allowed_values='base', 'root')

{'type_x': 'base', 'atr1': 'qwer'}
1 validation error for ParsingModel[typing_extensions.Annotated[Union[__main__.MyBase, __main__.MyRoot], FieldInfo(default=PydanticUndefined, discriminator='type_x', extra={})]]
__root__ -> MyBase -> atr2
  field required (type=value_error.missing)

{'type_x': 'root'}
2 validation errors for ParsingModel[typing_extensions.Annotated[Union[__main__.MyBase, __main__.MyRoot], FieldInfo(default=PydanticUndefined, discriminator='type_x', extra={})]]
__root__ -> MyRoot -> atr1
  field required (type=value_error.missing)
__root__ -> MyRoot -> atr3
  field required (type=value_error.missing)
→ Ссылка