在使用 Pydantic 进行数据验证和设置时,你可能会遇到一个问题:当多个类有共同字段时,Pydantic 可能无法正确地区分这些类。让我通过一个简单的例子来解释这个问题。

案例

我这里是 pydantic==1.10.11 环境,不确定是否在其他版本存在。

from pydantic import BaseModel

class A(BaseModel):
    a: str

class B(BaseModel):
    a: str
    b: int

class C(BaseModel):
    c: list[A | B]

c = C(c=[A(a="a1"), B(a="a2", b=2)])
print([type(x) for x in c.c])

输出结果为 [<class '__main__.A'>, <class '__main__.A'>],这意味着 B 类型的实例被错误地识别为 A 类型的实例。

解决方案

为了解决这个问题,我们可以使用 typing.Annotated 和 pydantic.Field 的 discriminator 参数。这允许我们指定一个字段作为区分不同类的标识符。

from typing import Annotated, Literal, Union
from pydantic import BaseModel, Field

class A(BaseModel):
    t: Literal["A"] = "A"
    a: str

class B(BaseModel):
    t: Literal["B"] = "B"
    a: str
    b: int

t = Annotated[Union[A, B], Field(discriminator="t")]

class C(BaseModel):
    c: list[t]

c = C(c=[A(a="a1"), B(a="a2", b=2)])
print([type(x) for x in c.c])

Field的 discriminator必须是所有可能类型的公共字段且为 Literal类型。

更简便的方式

如果你不想为每个类添加额外的标识字段,你可以考虑将共同字段提取到一个单独的基类中,然后让其他类继承这个基类。

from pydantic import BaseModel

class A(BaseModel):
    a: str

class B(A): # 继承有公共字段的基类
    a: str
    b: int

class C(BaseModel):
    c: list[A] # 只需要标识基类

c = C(c=[A(a="a1"), B(a="a2", b=2)])
print([type(x) for x in c.c])

输出 [<class '__main__.A'>, <class '__main__.B'>],这样,你就可以在不添加额外字段的情况下解决问题,同时也保持了代码的整洁和可读性。