如何在Python中为特定输入类型强制匹配对应输出类型？

问题背景 #

现有一组Pydantic动物模型，每个模型都有对应的"带帽子"子类：

from pydantic import BaseModel

class Horse(BaseModel):
    speed: str
    race_wins: int

class HorseWithHat(Horse):
    hat_color: str

class Snake(BaseModel):
    length: str
    poisonous: bool

class SnakeWithHat(Snake):
    hat_color: str

# 更多动物模型...

需要实现一个give_hat函数，要求：

• 输入某类动物实例，返回对应类型的"带帽子"实例（比如输入Horse必须返回HorseWithHat，不能返回SnakeWithHat）
• 尽量减少重复代码（比如避免每个带帽子模型都手动写hat_color字段）

解决方案一：用Mixin简化带帽子模型定义 #

首先用一个Mixin类统一定义hat_color字段，避免重复编写：

class WithHatMixin(BaseModel):
    hat_color: str

# 重新定义带帽子模型，只需继承基础动物+Mixin
class HorseWithHat(Horse, WithHatMixin):
    pass

class SnakeWithHat(Snake, WithHatMixin):
    pass

解决方案二：类型安全的give_hat函数实现 #

通过TypeVar绑定具体动物类型，配合类型映射字典，让类型检查器能正确推断输入输出的对应关系：

from typing import TypeVar, Type, Union

# 定义基础动物类型的Union
Animal = Union[Horse, Snake]
# 定义带帽子动物类型的Union
AnimalWithHat = Union[HorseWithHat, SnakeWithHat]

# 定义类型变量，绑定到每个具体的基础动物类型
T = TypeVar("T", Horse, Snake)

# 建立基础类型到带帽子类型的映射
TYPE_MAP: dict[Type[T], Type[AnimalWithHat]] = {
    Horse: HorseWithHat,
    Snake: SnakeWithHat,
    # 新增动物时只需在这里添加映射
}

def give_hat(animal: T, hat_color: str) -> TYPE_MAP[T]:
    # 将原动物的字段转成字典，加上hat_color后实例化对应带帽子模型
    return TYPE_MAP[type(animal)](**animal.model_dump(), hat_color=hat_color)

效果验证 #

此时类型检查器（如mypy、Pyright）会正确识别：

• 输入Horse实例，返回类型为HorseWithHat
• 输入Snake实例，返回类型为SnakeWithHat
• 不会出现输入Horse却返回SnakeWithHat的类型错误

进阶方案：用泛型模型自动生成带帽子类型 #

如果不想手动定义每个XWithHat模型，可以用Pydantic泛型模型实现动态组合字段：

from pydantic import BaseModel, ConfigDict
from typing import TypeVar, Generic, Union

AnimalT = TypeVar("AnimalT", bound=Union[Horse, Snake])

class WithHat(Generic[AnimalT], BaseModel):
    model_config = ConfigDict(from_attributes=True, extra="forbid")
    
    # 包含基础动物的所有字段
    animal: AnimalT
    hat_color: str

# 使用示例
horse = Horse(speed="fast", race_wins=5)
horse_with_hat = WithHat(animal=horse, hat_color="red")

这种方式的字段是嵌套结构（animal字段包含原动物数据），如果需要平级字段，可以结合Pydantic的model_computed_fields或者自定义字段生成逻辑，但嵌套结构在多数场景下已经够用。

内容的提问来源于stack exchange，提问作者Yot Yot5