如何在Pydantic中自定义带固定鉴别器的Union类型注解？

在Pydantic中，我们可以通过Annotated结合Union和Field(discriminator="disc")实现带鉴别器的联合类型，示例代码如下：

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

class A(BaseModel):
    disc: Literal["a"]
    other: str = "asdf"

class B(BaseModel):
    disc: Literal["b"]
    other: int = 23

class Thing(BaseModel):
    foo: Annotated[Union[A, B], Field(discriminator="disc")]

现在需要创建CustomUnion[A,B,...]类型，使其等价于Annotated[Union[A,B,...], Field(discriminator="disc")]，并且像原生Union一样强制要求至少传入两个类型参数。以下是两种可行的实现方案：

方案一：Python 3.11+ 原生支持（推荐） #

利用Python 3.11引入的TypeVarTuple和Unpack特性，可以更优雅地处理可变数量的类型参数：

from typing import TypeVar, Union, TypeVarTuple, Unpack
from pydantic import BaseModel, Field, Annotated

# 定义绑定BaseModel的类型变量和类型元组变量
BaseModelT = TypeVar("BaseModelT", bound=BaseModel)
ModelsT = TypeVarTuple("ModelsT", bound=BaseModel)

class CustomUnionMeta(type):
    def __class_getitem__(cls, items: tuple[Unpack[ModelsT]]) -> Annotated[Union[Unpack[ModelsT]], Field]:
        # 强制要求至少两个类型参数
        if len(items) < 2:
            raise TypeError("CustomUnion requires at least two type arguments")
        # 返回等价的Annotated联合类型
        return Annotated[Union[items], Field(discriminator="disc")]

class CustomUnion(metaclass=CustomUnionMeta):
    pass

# 使用示例
class A(BaseModel):
    disc: Literal["a"]
    other: str = "asdf"

class B(BaseModel):
    disc: Literal["b"]
    other: int = 23

class C(BaseModel):
    disc: Literal["c"]
    value: float = 3.14

class Thing(BaseModel):
    foo: CustomUnion[A, B]  # 等价于Annotated[Union[A,B], Field(discriminator="disc")]
    bar: CustomUnion[A, B, C]  # 支持多类型联合

方案二：兼容Python 3.10及以下版本 #

如果需要兼容更早的Python版本，可以通过元类的__class_getitem__手动处理参数：

from typing import TypeVar, Union, Tuple
from pydantic import BaseModel, Field, Annotated

BaseModelT = TypeVar("BaseModelT", bound=BaseModel)

class CustomUnionMeta(type):
    def __class_getitem__(cls, items: Tuple[BaseModelT, ...]) -> Annotated[Union[BaseModelT, ...], Field]:
        # 处理单个参数的情况（转为元组）
        if not isinstance(items, tuple):
            items = (items,)
        # 强制要求至少两个类型参数
        if len(items) < 2:
            raise TypeError("CustomUnion requires at least two type arguments")
        # 返回等价的Annotated联合类型
        return Annotated[Union[items], Field(discriminator="disc")]

class CustomUnion(metaclass=CustomUnionMeta):
    pass

# 使用方式与方案一完全一致

两种方案都实现了简洁的调用形式，同时确保了CustomUnion必须接收至少两个类型参数，完全满足需求。

内容的提问来源于stack exchange，提问作者Ivan GJ