如何在mypy中定义支持子类兼容的泛型引用类型，适配Pydantic场景？

解决方案 #

核心思路是将MyBaseModel改为泛型基类，让id字段的类型自动绑定到当前子类的类型，修改后的完整代码如下：

from typing import TypeVar, Generic
from pydantic import BaseModel, Field

ModelType = TypeVar("ModelType", bound="MyBaseModel")
# 新增用来表示当前子类类型的类型变量
SelfModelT = TypeVar("SelfModelT", bound="MyBaseModel")


class IDType(str):
    pass


class ReferenceId(IDType, Generic[ModelType]):
    pass


# 改为泛型基类
class MyBaseModel(BaseModel, Generic[SelfModelT]):
    # id类型绑定为当前子类的ReferenceId
    id: ReferenceId[SelfModelT] = Field(default_factory=lambda: ReferenceId[SelfModelT]())


class MySubclassModel(MyBaseModel["MySubclassModel"]):
    pass


class AnotherSubclassModel(MyBaseModel["AnotherSubclassModel"]):
    reference_to_my_subclass: ReferenceId[MySubclassModel]


class YetAnotherSubclassModel(MyBaseModel["YetAnotherSubclassModel"]):
    reference_to_another_subclass: ReferenceId[AnotherSubclassModel]


# 原示例此处传入的是MyBaseModel的id，本身就和要求的ReferenceId[MySubclassModel]类型不匹配，属于正确报错
# 修改为传入MySubclassModel的id即可正常通过校验
a = AnotherSubclassModel(reference_to_my_subclass=MySubclassModel().id)
# 此处现在可以正常通过校验，不会报错
x = YetAnotherSubclassModel(reference_to_another_subclass=a.id)


# 原有符合预期的逻辑保持正常
def my_function(reference_to_my_subclass: ReferenceId[MySubclassModel]):
    pass

my_function(a.reference_to_my_subclass)

# 错误传参依然会被拦截
my_function(a.id) # 正确报错：类型不匹配
my_function(x.reference_to_another_subclass) # 正确报错：类型不匹配

改动说明 #

• 新增SelfModelT类型变量，用来标记MyBaseModel子类的自身类型
• 将MyBaseModel改为泛型类，每个子类继承时需要传入自身作为泛型参数，这样子类的id字段会自动被标记为ReferenceId[当前子类]类型
• 调整id的默认工厂为lambda表达式，动态生成对应泛型的ReferenceId实例，运行时依然是字符串子类，完全兼容原有的MongoDB存储逻辑

如果你使用Python 3.11+和Pydantic 2.0+，可以用typing.Self简化写法，不需要每个子类显式传入自身作为泛型参数：

from typing import Self

class MyBaseModel(BaseModel):
    id: ReferenceId[Self] = Field(default_factory=lambda: ReferenceId[Self]())

class MySubclassModel(MyBaseModel):
    pass

内容的提问来源于stack exchange，提问作者Guilherme Caminha