将pydantic基模型属性强制转换为列表

Question

我有一个继承自基模型的类，定义如下:

class MessageInfo(CustomBaseModel):
completion_code: int = None
code: str = None
message: str = None
severity: str = None

class OrderResponseMessageInfo(CustomBaseModel):
message_info: Union[List[MessageInfo], MessageInfo] = None

希望此message_info能够使用列表或字典，但在将其序列化为字典或json时始终如一地将其作为列表生成。

由于这是API响应的解析器，并且API响应不一致，即有时消息信息块是像这样的数组:

{
"MessageInfo": [
{
"CompletionCode": 517,
"Code": "1B61",
"Message": "some message",
"Severity": "Warning"
},
{
"CompletionCode": 9886,
"Code": "1B61",
"Message": "Another message",
"Severity": "Warning"
}
]
}

有时是像这样的字典

{
"MessageInfo": {
"CompletionCode": 517,
"Code": "1B61",
"Message": "some message",
"Severity": "Warning"
}
}

当消息信息是一个字典，我做OrderResponseMessageInfo.dict()，我得到

{
"OrderLineMessageInfo": {
"MessageInfo": {
"CompletionCode": 517,
"Code": "1B61",
"Message": "Routing completed. Confirmation required.",
"Severity": "Warning"
}
}
}

我想要的是，如果消息信息是一个字典，那么我应该能够将其类型转换为列表，这样当我做OrderResponseMessageInfo.dict()时，我应该得到输出:

{
"OrderLineMessageInfo": {
"MessageInfo": [
{
"CompletionCode": 517,
"Code": "1B61",
"Message": "Routing completed. Confirmation required.",
"Severity": "Warning"
}
]
}
}

有办法做到这一点吗?

Answer 1

作为一般规则，您应该根据实际需要的模式来定义您的模型，而不是根据可能得到的。

如果您希望一个字段是列表类型，那么就这样定义它。如果您希望在解析/初始化期间允许将单个元素转换为单项列表作为特殊情况，您可以定义一个自定义的pre=True字段验证器来完成此操作。

无论哪种方式，在解析和验证过程结束时，您应该最终得到所需的Schema，而不是你需要在某个阶段再次打乱的东西。

的例子:

from collections.abc import Mapping
from pydantic import BaseModel, validator

class Foo(BaseModel):
x: int
y: str

class Bar(BaseModel):
foos: list[Foo]
@validator("foos", pre=True)
def single_to_list(cls, v: object) -> object:
if isinstance(v, (Mapping, Foo)):
return [Foo.parse_obj(v)]
return v

演示:

b1 = Bar.parse_raw('{"foos": [{"x": 1, "y": "spam"}, {"x": -1, "y": "eggs"}]}')
b2 = Bar.parse_raw('{"foos": {"x": 42, "y": "ham"}}')
f = Foo(x=69, y="beans")
b3 = Bar.parse_obj({"foos": f})
b4 = Bar.parse_obj({"foos": [f, f]})
print(b1.json(indent=4))
print(b2.json(indent=4))
print(b3.json(indent=4))
print(b4.json(indent=4))

输出:

{
"foos": [
{
"x": 1,
"y": "spam"
},
{
"x": -1,
"y": "eggs"
}
]
}

{
"foos": [
{
"x": 42,
"y": "ham"
}
]
}

{
"foos": [
{
"x": 69,
"y": "beans"
}
]
}

{
"foos": [
{
"x": 69,
"y": "beans"
},
{
"x": 69,
"y": "beans"
}
]
}

如果你想让message_info是一个列表，那么定义它为list[MessageInfo]，并编写一个自定义验证器来尝试强制任何映射或MessageInfo实例到包含一个MessageInfo实例的列表。

当然，您可以重写转储/序列化方法(如.dict和.json)，但我不建议这样做。1)这会导致更多的代码，2)它会降低效率，3)模型模式仍然会误导，因为您将总是将该字段转储为列表，但是模式会暗示单个实例可能会出现。