根据初始化参数初始化两个 Pydantic 模型之一

Question

我有一个MessageModel版本号为 的Pydantic 模型类Literal。现在我们的需求发生了变化，我们需要另一个MessageModel具有更高版本号的版本，因为 的属性MessageModel已经发生了变化。我想要一个类，我可以在其中将版本号作为构造函数的参数。有人有想法吗？

以下是型号：

from typing import Literal
from pydantic import BaseModel


class MessageModelV1(BaseModel):
    version: Literal[1]
    bar: str
        
class MessageModelV2(BaseModel):
    version: Literal[2]
    foo: str

我想要的是一个初始化正确MessageModel版本的类：

model = MessageModel(version=2, ...)

Answer 1

受歧视工会

您可以定义两个或多个 Pydantic 模型的可辨别联合。这将允许您根据数据中提供的鉴别器字段的值实例化正确的模型。

您可以采取几种略有不同的方法。

选项 A：带注释的类型别名

您无需定义“组合”现有模型的新模型，而是为这些模型的并集定义一个类型别名，并用于typing.Annotated添加鉴别器信息。

# Pydantic v1

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


class MessageModelV1(BaseModel):
    version: Literal[1]
    bar: str


class MessageModelV2(BaseModel):
    version: Literal[2]
    foo: str


MessageModel = Annotated[
    Union[MessageModelV1, MessageModelV2],
    Field(discriminator="version"),
]


data1 = {"version": 1, "bar": "a"}
data2 = {"version": 2, "foo": "b"}
obj1 = parse_obj_as(MessageModel, data1)
obj2 = parse_obj_as(MessageModel, data2)
print(obj1, type(obj1))  # version=1 bar='a' <class '__main__.MessageModelV1'>
print(obj2, type(obj2))  # version=2 foo='b' <class '__main__.MessageModelV2'>

优点

• 最少的附加代码，只需一个类型定义。
• 您将获得原始模型之一的实际实例。
• 静态类型检查器可以将类型推断为原始模型的联合。（这使得 IDE 能够提供一些有用的自动建议。）

缺点

• 您不能直接实例化，MessageModel因为它是类型构造而不是模型类。
• 要构造实例，您必须使用附加parse_obj_as函数。
• 无法静态确定解析后会出现哪个确切模型，因为这取决于数据。

更新：Pydantic v2

非常相似，但使用TypeAdapter代替parse_as_obj：

# Pydantic v2

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


class MessageModelV1(BaseModel):
    version: Literal[1]
    bar: str


class MessageModelV2(BaseModel):
    version: Literal[2]
    foo: str


MessageModel = TypeAdapter(Annotated[
    Union[MessageModelV1, MessageModelV2],
    Field(discriminator="version"),
])


data1 = {"version": 1, "bar": "a"}
data2 = {"version": 2, "foo": "b"}
obj1 = MessageModel.validate_python(data1)
obj2 = MessageModel.validate_python(data2)
print(obj1, type(obj1))  # version=1 bar='a' <class '__main__.MessageModelV1'>
print(obj2, type(obj2))  # version=2 foo='b' <class '__main__.MessageModelV2'>

不幸的是，到目前为止，Mypy 似乎无法正确推断结果模型实例的类型，而 Pyright 可以。

选项 B：自定义根类型

您定义一个新模型并将其__root__类型设置为原始模型之间的判别并集。

然后，您可以将其自定义到您认为合适的程度，以使它的实例“感觉”像任何原始的底层模型。

在下面的示例中，我重写了该__init__方法，以便您可以像问题中所述那样初始化它。我还通过底层根模型使其可迭代__str__，并添加了自定义和__repr__方法，以便实例显示为底层根类型。

# Pydantic v1

from collections.abc import Iterator
from typing import Any, Literal, Union
from pydantic import BaseModel, Field


class MessageModelV1(BaseModel):
    version: Literal[1]
    bar: str


class MessageModelV2(BaseModel):
    version: Literal[2]
    foo: str


MessageType = Union[MessageModelV1, MessageModelV2]


class MessageModel(BaseModel):
    __root__: MessageType = Field(discriminator="version")

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(__root__=kwargs)

    def __iter__(self) -> Iterator[tuple[str, Any]]:  # type: ignore[override]
        yield from self.__root__

    def __str__(self) -> str:
        return str(self.__root__)

    def __repr__(self) -> str:
        return repr(self.__root__)

使用相同的演示数据，您现在得到略有不同的结果：

# Pydantic v1

obj1 = MessageModel(version=1, bar="a")
obj2 = MessageModel(version=2, foo="b")
print(obj1, type(obj1))  # version=1 bar='a' <class '__main__.MessageModel'>
print(obj2, type(obj2))  # version=2 foo='b' <class '__main__.MessageModel'>

print(hasattr(obj1, "__root__"))  # True
print(hasattr(obj1, "bar"))       # False
print(hasattr(obj1, "version"))   # False

您会注意到，对象的类型现在当然是新的MessageModel，而不是底层原始模型之一。

此外，如果不进行更多自定义，实例将仅具有__root__指向底层模型实例的字段，而不会具有该模型的实际字段。如果您想将属性访问传递给根模型，则必须覆盖__getattr__/ __setattr__：

# Pydantic v1

...

class MessageModel(BaseModel):
    __root__: MessageType = Field(discriminator="version")

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(__root__=kwargs)

    def __getattr__(self, name: str) -> Any:
        if name == "__root__":
            return self.__root__
        return getattr(self.__root__, name)

    def __setattr__(self, name: str, value: Any) -> None:
        if name == "__root__":
            self.__root__ = value
        setattr(self.__root__, name, value)

    ...

# Pydantic v1

obj1 = MessageModel(version=1, bar="a")

print(obj1.version)  # 1
print(obj1.bar)      # a

优点

• 直接实例化是可能的，因为您正在处理实际的模型类。
• 许多定制和额外验证的选项。

缺点

• 如果没有任何自定义，您将始终处理“代理”模型，其__root__字段将指向“实际”模型。
• 让外部模型感觉/表现得像底层模型之一需要大量额外的样板。您可能可以将所有/大部分内容分解到一个单独的混合类中，但仍然需要编写代码。
• 您可能会“忘记”您正在处理代理模型，而不是根模型，然后遇到意外的结果。（例如，尝试调用obj1.dict()最后一个示例并查看输出。）
• 静态分析推断的类型始终是MessageModel.

聚苯乙烯

您可以稍微绕过以下事实：外部模型MessageModel实际上并不是底层模型的子类型typing.TYPE_CHECKING。例子：

# Pydantic v1

from typing import Any, Literal, Union, TYPE_CHECKING
from pydantic import BaseModel, Field


class MessageModelV1(BaseModel):
    version: Literal[1]
    bar: str


class MessageModelV2(BaseModel):
    version: Literal[2]
    foo: str


MessageType = Union[MessageModelV1, MessageModelV2]

if TYPE_CHECKING:
    class MessageModel(MessageModelV1, MessageModelV2):
        pass
else:
    class MessageModel(BaseModel):
        __root__: MessageType = Field(discriminator="version")

        def __init__(self, **kwargs: Any) -> None:
            super().__init__(__root__=kwargs)

        ...

obj1 = MessageModel(version=1, bar="a")

从静态分析的角度来看，MessageModel它现在是其他两种的子类型，这意味着您的 IDE 将为您提供相应的见解。例如，输入obj1.PyCharm 将建议属性foo和bar。从打字角度来看，这使得它本质上等同于Union别名方法（选项 A）。

但这当然是一个谎言，因为实际的运行时 MessageModel不是这两个的子类。但是，如果您确切地知道需要什么接口才能使其“感觉”像子类，则可以使其工作。

我不建议在混淆方面走那么远。

更新：Pydantic v2

非常相似，但使用RootModel：

# Pydantic v2

from collections.abc import Iterator
from typing import Any, Literal, Union
from pydantic import BaseModel, Field, RootModel


class MessageModelV1(BaseModel):
    version: Literal[1]
    bar: str


class MessageModelV2(BaseModel):
    version: Literal[2]
    foo: str


MessageType = Union[MessageModelV1, MessageModelV2]


class MessageModel(RootModel):
    root: MessageType = Field(discriminator="version")

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(root=kwargs)

    def __iter__(self) -> Iterator[tuple[str, Any]]:  # type: ignore[override]
        yield from self.root

    def __str__(self) -> str:
        return str(self.root)

    def __repr__(self) -> str:
        return repr(self.root)


obj1 = MessageModel(version=1, bar="a")
obj2 = MessageModel(version=2, foo="b")
print(obj1, type(obj1))  # version=1 bar='a' <class '__main__.MessageModel'>
print(obj2, type(obj2))  # version=2 foo='b' <class '__main__.MessageModel'>

print(hasattr(obj1, "root"))     # True
print(hasattr(obj1, "bar"))      # False
print(hasattr(obj1, "version"))  # False

选项C：联合字段+代理构造函数

本质上基于选项 B，但不必经历为嵌套模型定义代理接口的所有麻烦，您只需编写一个看起来像类的自定义构造函数，并使其实例化外部模型，但仅返回其__root__价值：

# Pydantic v1

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


class MessageModelV1(BaseModel):
    version: Literal[1]
    bar: str


class MessageModelV2(BaseModel):
    version: Literal[2]
    foo: str


MessageType = Union[MessageModelV1, MessageModelV2]


class OuterMessageModel(BaseModel):
    __root__: MessageType = Field(discriminator="version")


def MessageModel(**kwargs: Any) -> MessageType:
    return OuterMessageModel.parse_obj(kwargs).__root__


obj1 = MessageModel(version=1, bar="a")
obj2 = MessageModel(version=2, foo="b")
print(obj1, type(obj1))  # version=1 bar='a' <class '__main__.MessageModelV1'>
print(obj2, type(obj2))  # version=2 foo='b' <class '__main__.MessageModelV2'>

优点

• 相当简单；代码不多。
• 直接实例化是可能的（尽管通过代理构造函数）。
• 您将获得原始模型之一的实际实例。
• 静态类型检查器可以将类型推断为原始模型的联合。（这使得 IDE 能够提供一些有用的自动建议。）

缺点

• 构造函数看起来像一个“假”类。（尽管这是一种相当常见的方法；例如，请参阅 PydanticField函数。）
• 仍然无法静态地确定解析后会出现哪个确切的模型，因为这取决于数据。

更新：Pydantic v2

非常相似，但使用RootModel：

# Pydantic v2

from typing import Any, Literal, Union
from pydantic import BaseModel, Field, RootModel


class MessageModelV1(BaseModel):
    version: Literal[1]
    bar: str


class MessageModelV2(BaseModel):
    version: Literal[2]
    foo: str


MessageType = Union[MessageModelV1, MessageModelV2]


class OuterMessageModel(RootModel):
    root: MessageType = Field(discriminator="version")


def MessageModel(**kwargs: Any) -> MessageType:
    return OuterMessageModel.model_validate(kwargs).root


obj1 = MessageModel(version=1, bar="a")
obj2 = MessageModel(version=2, foo="b")
print(obj1, type(obj1))  # version=1 bar='a' <class '__main__.MessageModelV1'>
print(obj2, type(obj2))  # version=2 foo='b' <class '__main__.MessageModelV2'>