Pydantic模型中动态子类联合类型的优雅实现：判别式联合与自动化策略

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

class Pet(BaseModel):
    """动物基类"""
    name: str
    age: int

class Dog(Pet):
    """狗类模型"""
    # 'type' 字段作为判别器，其值必须是 Literal["dog"]
    type: Literal["dog"] = "dog"
    breed: str

class Cat(Pet):
    """猫类模型"""
    # 'type' 字段作为判别器，其值必须是 Literal["cat"]
    type: Literal["cat"] = "cat"
    breed: str

# 定义判别式联合类型 AnyPet
# Annotated 用于添加元数据，Field(discriminator="type") 指定 'type' 字段为判别器
AnyPet = Annotated[Union[Dog, Cat], Field(discriminator="type")]

class Home(BaseModel):
    """家模型，包含一个宠物"""
    pet: AnyPet

# 示例数据
data = {
    "pet": {
        "type": "dog",  # 根据 "type" 字段的值，Pydantic 会自动解析为 Dog 实例
        "name": "Buddy",
        "age": 4,
        "breed": "Golden Retriever"
    }
}

# 创建 Home 实例并验证
home = Home(**data)
print(home)
# 输出: pet=Dog(name='Buddy', age=4, type='dog', breed='Golden Retriever')

data_cat = {
    "pet": {
        "type": "cat",
        "name": "Whiskers",
        "age": 2,
        "breed": "Siamese"
    }
}
home_cat = Home(**data_cat)
print(home_cat)
# 输出: pet=Cat(name='Whiskers', age=2, type='cat', breed='Siamese')

# 项目结构示例
your_project/
├── models/
│   ├── __init__.py  # 定义 AnyPet 及其所有子类
│   ├── pets.py      # 也可以将 Pet 基类和通用逻辑放在这里
│   ├── dogs.py      # 定义 Dog
│   └── cats.py      # 定义 Cat
└── main.py

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

# 假设 Pet、Dog、Cat 等类已在适当位置定义和导入
# 为了演示，我们再次定义它们
class Pet(BaseModel):
    name: str
    age: int

class Dog(Pet):
    type: Literal["dog"] = "dog"
    breed: str

class Cat(Pet):
    type: Literal["cat"] = "cat"
    breed: str

# 动态发现 Pet 的所有子类
valid_sub_classes = []
for sub_class in Pet.__subclasses__():
    # 验证子类是否包含判别器字段
    # Pydantic v2 使用 model_fields
    if "type" not in sub_class.model_fields:
        raise ValueError(f"子类 {sub_class.__name__} 缺少判别器 'type' 字段")
    # 进一步验证 'type' 字段是否为 Literal
    field_info = sub_class.model_fields["type"].annotation
    if not (hasattr(field_info, '__origin__') and field_info.__origin__ is Literal):
         raise ValueError(f"子类 {sub_class.__name__} 的 'type' 字段必须是 Literal 类型")

    valid_sub_classes.append(sub_class)

# 使用动态发现的子类列表创建判别式联合
if not valid_sub_classes:
    # 处理没有子类的情况，例如定义一个默认的 AnyPet
    AnyPet = Annotated[Pet, Field(discriminator="type")] # 或者根据实际需求处理
else:
    AnyPet = Annotated[Union[tuple(valid_sub_classes)], Field(discriminator="type")]

print("动态生成的 AnyPet 类型:", AnyPet)

class Home(BaseModel):
    pet: AnyPet

# 再次测试
data = {
    "pet": {
        "type": "dog",
        "name": "Buddy",
        "age": 4,
        "breed": "Golden Retriever"
    }
}
home = Home(**data)
print(home)

# my_module.py
from pydantic import BaseModel, Field
from typing import Literal, Annotated, Union

# 假设 Pet 类在这里定义
class Pet(BaseModel):
    name: str
    age: int

# 其他模块可能定义了 Dog 和 Cat
# ...

def get_any_pet_type() -> Annotated[Union, Field]:
    """
    动态生成并返回 AnyPet 判别式联合类型。
    此函数应在所有 Pet 的子类模块都已导入后调用。
    """
    valid_sub_classes = []
    for sub_class in Pet.__subclasses__():
        if "type" not in sub_class.model_fields:
            raise ValueError(f"子类 {sub_class.__name__} 缺少判别器 'type' 字段")
        valid_sub_classes.append(sub_class)

    if not valid_sub_classes:
        # 如果没有发现子类，返回一个默认的类型或抛出错误
        return Annotated[Pet, Field(discriminator="type")] 

    return Annotated[Union[tuple(valid_sub_classes)], Field(discriminator="type")]

# main.py
from pydantic import BaseModel
from my_module import get_any_pet_type # 导入获取联合类型的函数

# 假设其他模块（如 dogs.py, cats.py）已被导入，定义了 Dog 和 Cat
# from .other_modules import Dog, Cat # 实际项目中会这样导入

# 示例：模拟 Dog 和 Cat 在其他地方被定义
class Dog(Pet): # Pet 假设在 my_module.py 中
    type: Literal["dog"] = "dog"
    breed: str

class Cat(Pet):
    type: Literal["cat"] = "cat"
    breed: str

# 在所有子类都已加载后，调用函数获取 AnyPet 类型
AnyPet = get_any_pet_type()

class Home(BaseModel):
    """Home class"""
    pet: AnyPet

# 测试
data = {
    "pet": {
        "type": "cat",
        "name": "Luna",
        "age": 1,
        "breed": "Persian"
    }
}
home = Home(**data)
print(home)