在使用 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'>],这样,你就可以在不添加额外字段的情况下解决问题,同时也保持了代码的整洁和可读性。