Pydantic проверить, что в строке только допустимые значения (StrEnum)

Question

Сильно упрощенный json

{
    "tariffOptions": "qwertyTest,qwerty,test"
}

qwertyTest, qwerty, test - это все возможные значения, располагаться они могут в произвольном порядке

Могут быть все, одно\два из них или поле может отсутствовать

Нужно проверить, что в строке только допустимые значения

Пробовал решить так, но это не работает

class TariffOptions(StrEnum):
    qwerty_test = 'qwertyTest'
    qwerty = 'qwerty'
    test = 'test'

class Data(Base):
    tariff_options: Optional[TariffOptions] = None

Answer 1

Тут у вас по сути в tariff_options должен получиться список объектов TariffOptions, а не Optional. Нужно исходные данные подготовить, разбить по запятым на список, потом это уже распарсится как список enum.

Разбивать строку можно внутри метода с декоратором @model_validator(mode='before') (см. Model validators) - в него приходит json в виде словаря, можно этот словарь подправить и вернуть из этого метода:

from enum import StrEnum
from typing import Any
from pydantic import BaseModel, model_validator


class TariffOptions(StrEnum):
    qwerty_test = 'qwertyTest'
    qwerty = 'qwerty'
    test = 'test'


class Data(BaseModel):
    tariff_options: list[TariffOptions]
    
    @model_validator(mode="before")
    @classmethod
    def from_literal(cls, data: Any) -> Any:
        # Тут одновременно с разбиением происходит изменение
        # имени поля в snake_case
        # - по хорошему это нужно делать отдельно через алиасы
        data["tariff_options"] = data["tariffOptions"].split(",")
        return data


source_json = """
{
    "tariffOptions": "qwertyTest,qwerty,test"
}
"""

print(repr(Data.model_validate_json(source_json)))

Вывод:

Data(tariff_options=[<TariffOptions.qwerty_test: 'qwertyTest'>, <TariffOptions.qwerty: 'qwerty'>, <TariffOptions.test: 'test'>])

Если добавить что-то, чего в enum нет, например:

{
    "tariffOptions": "qwertyTest,erhrthrt"
}

То будет такая ошибка:

Traceback (most recent call last):
  File "/home/insolor/Projects/stackoverflow/test_pydantic/test.py", line 27, in <module>
    print(repr(Data.model_validate_json(source_json)))
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/insolor/Projects/stackoverflow/test_pydantic/venv/lib/python3.12/site-packages/pydantic/main.py", line 622, in model_validate_json
    return cls.__pydantic_validator__.validate_json(json_data, strict=strict, context=context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
pydantic_core._pydantic_core.ValidationError: 1 validation error for Data
tariff_options.1
  Input should be 'qwertyTest', 'qwerty' or 'test' [type=enum, input_value='erhrthrt', input_type=str]
    For further information visit https://errors.pydantic.dev/2.9/v/enum

БЛОГ НА HUSL

Pydantic проверить, что в строке только допустимые значения (StrEnum)

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