如何在Pydantic v2验证错误中移除loc中的鉴别器值？

问题场景 #

使用Pydantic v2的鉴别器（discriminator）时，验证错误的loc字段会包含鉴别器对应的值（比如示例中的'shark'），导致loc输出为('shark', 'ferocity_scale')，需要移除这个多余的鉴别器值。

示例代码：

from pprint import pprint
from typing import Literal, Union, Annotated, Any

from pydantic import BaseModel, Field, RootModel, ValidationError
from pydantic.main import Model


class Tiger(BaseModel):
    animal_type: Literal["tiger"] = "tiger"
    ferocity_scale: float = Field(..., ge=0, le=10)


class Shark(BaseModel):
    animal_type: Literal["shark"] = "shark"
    ferocity_scale: float = Field(..., ge=0, le=10)


class Lion(BaseModel):
    animal_type: Literal["lion"] = "lion"
    ferocity_scale: float


class WildAnimal(RootModel):
    root: Annotated[Union[Tiger, Shark, Lion], Field(..., discriminator='animal_type')]


try:
    my_shark = WildAnimal.model_validate({'animal_type': 'shark', 'ferocity_scale': 115})
except ValidationError as exc:
    pprint(exc.errors())

运行后错误输出：

[{'ctx': {'le': 10.0},
  'input': 115,
  'loc': ('shark', 'ferocity_scale'),  # 需要移除此处的"shark"
  'msg': 'Input should be less than or equal to 10',
  'type': 'less_than_equal'}]

解决方案 #

方法一：重写model_validate直接修改错误信息 #

通过重写WildAnimal的model_validate方法，捕获ValidationError后修改错误列表中的loc字段，再重新抛出修改后的错误：

from pprint import pprint
from typing import Literal, Union, Annotated, Any

from pydantic import BaseModel, Field, RootModel, ValidationError
from pydantic.main import Model


class Tiger(BaseModel):
    animal_type: Literal["tiger"] = "tiger"
    ferocity_scale: float = Field(..., ge=0, le=10)


class Shark(BaseModel):
    animal_type: Literal["shark"] = "shark"
    ferocity_scale: float = Field(..., ge=0, le=10)


class Lion(BaseModel):
    animal_type: Literal["lion"] = "lion"
    ferocity_scale: float


class WildAnimal(RootModel):
    root: Annotated[Union[Tiger, Shark, Lion], Field(..., discriminator='animal_type')]

    @classmethod
    def model_validate(
        cls: type[Model],
        obj: Any,
        *,
        strict: bool | None = None,
        from_attributes: bool | None = None,
        context: dict[str, Any] | None = None
    ) -> Model:
        try:
            return super().model_validate(obj, strict=strict, from_attributes=from_attributes, context=context)
        except ValidationError as exc:
            modified_errors = []
            # 遍历错误列表，移除loc中的鉴别器值
            for err in exc.errors():
                if len(err['loc']) >= 1 and err['loc'][0] in {'tiger', 'shark', 'lion'}:
                    err['loc'] = err['loc'][1:]
                modified_errors.append(err)
            # 重新构造并抛出ValidationError
            raise ValidationError.from_exception_data(title=exc.title, errors=modified_errors) from exc


try:
    my_shark = WildAnimal.model_validate({'animal_type': 'shark', 'ferocity_scale': 115})
except ValidationError as exc:
    pprint(exc.errors())

运行后错误的loc会变为('ferocity_scale',)，符合需求。

方法二：动态识别鉴别器值（通用型） #

如果不想硬编码鉴别器的可能取值，可以通过模型字段信息动态获取所有鉴别器值，适配后续新增的模型：

@classmethod
def model_validate(
    cls: type[Model],
    obj: Any,
    *,
    strict: bool | None = None,
    from_attributes: bool | None = None,
    context: dict[str, Any] | None = None
) -> Model:
    try:
        return super().model_validate(obj, strict=strict, from_attributes=from_attributes, context=context)
    except ValidationError as exc:
        # 动态获取所有鉴别器的可能取值
        discriminator_values = set()
        union_type = cls.model_fields['root'].annotation.__args__
        for model in union_type:
            if 'animal_type' in model.model_fields:
                val = model.model_fields['animal_type'].default
                discriminator_values.add(val)
        
        modified_errors = []
        for err in exc.errors():
            if len(err['loc']) >= 1 and err['loc'][0] in discriminator_values:
                err['loc'] = err['loc'][1:]
            modified_errors.append(err)
        
        raise ValidationError.from_exception_data(title=exc.title, errors=modified_errors) from exc

这种方式无需硬编码具体的鉴别器值，新增动物模型时会自动识别处理。

内容的提问来源于stack exchange，提问作者Альберт Александров