我有这样的模型:
class Foo(BaseModel):
protocol: str
protocol_params: Union[ProtocolOneParam, ProtocolTwoParam]
ProtocolOneParam
和
ProtocolTwoParam
没有具有可区分值的相同字段,因此我可以将它们用作
Discriminator
,而我可以理解哪个模型应该用于
protocol_params
的唯一方法是通过值
protocol
可以是
"protocol1"
或
"protocol2"
if
protocol == "protocol1"
则
protocol_params
应该使用
ProtocolOneParam
模型进行反序列化等等。
在实际用例中,有超过 5 个协议,其中没有任何可区分的字段。
pydantic 有什么方法可以实现我所需要的吗?
当然,Pydantic 可以实现想要的功能。可以使用
@validator
和
root_validator
来根据
protocol
字段的值动态设置
protocol_params
的类型。
from typing import Union, Literal
from pydantic import BaseModel, validator
class ProtocolOneParam(BaseModel):
param1: str
class ProtocolTwoParam(BaseModel):
param2: int
class Foo(BaseModel):
protocol: Literal["protocol1", "protocol2"]
protocol_params: Union[ProtocolOneParam, ProtocolTwoParam]
@validator('protocol_params', pre=True)
def validate_protocol_params(cls, value, values):
if values['protocol'] == "protocol1":
return ProtocolOneParam(**value)
elif values['protocol'] == "protocol2":
return ProtocolTwoParam(**value)
else:
raise ValueError(f"Invalid protocol: {values['protocol']}")
# 示例
foo1 = Foo(protocol="protocol1", protocol_params={"param1": "hello"})
foo2 = Foo(protocol="protocol2", protocol_params={"param2": 123})
print(foo1) # protocol='protocol1', protocol_params=ProtocolOneParam(param1='hello')
print(foo2) # protocol='protocol2', protocol_params=ProtocolTwoParam(param2=123)
解释:
-
使用
Literal
类型限制protocol
字段的取值范围,使其只能是"protocol1"
或"protocol2"
。 -
@validator('protocol_params', pre=True)
装饰器定义了一个验证器,它会在protocol_params
字段被赋值之前运行。 -
在验证器中,我们根据
values['protocol']
的值,将value
参数(即原始数据)转换成对应的ProtocolOneParam
或ProtocolTwoParam
对象。 -
如果
protocol
的值不匹配,则抛出ValueError
异常。
对于拥有更多协议的情况,可以扩展
validate_protocol_params
验证器,根据
protocol
的值选择对应的模型。
这种方法可以有效地解决的问题,并且易于扩展和维护。
标签:python,python-3.x,pydantic,pydantic-v2 From: 77887248