如何为返回关联类型的函数标注类型以适配子类？

问题场景 #

父类Foo包含返回Bar类型的方法，子类MyFoo的对应方法返回MyBar（Bar的子类）。运行时代码能正确返回对应类型，但mypy无法识别MyFoo.to_bar()的实际返回类型为MyBar，且不想通过逐个重写父类方法的方式解决。

原始代码 #

from __future__ import annotations

class Foo:
    def __init__(self, value):
        self.value = value

    @property
    def _bar(self) -> type[Bar]:
        return Bar

    def to_bar(self) -> Bar:
        return self._bar(self.value)

class Bar:
    def __init__(self, value):
        self.value = value

class MyFoo(Foo):
    @property
    def _bar(self) -> type[MyBar]:
        return MyBar

class MyBar(Bar):
    ...

print(type(Foo(3).to_bar()))
print(type(MyFoo(3).to_bar()))

运行输出 #

$ python t.py
<class '__main__.Bar'>
<class '__main__.MyBar'>

mypy检测结果 #

添加类型检测代码：

reveal_type(Foo(3).to_bar())
reveal_type(MyFoo(3).to_bar())

执行mypy后输出：

$ mypy t.py
t.py:28: note: Revealed type is "t.Bar"
t.py:29: note: Revealed type is "t.Bar"
Success: no issues found in 1 source file

解决方案：使用泛型关联父类与返回类型 #

通过typing.Generic和类型变量，让父类Foo成为泛型类，将返回类型与子类关联，无需重写每个方法即可让mypy正确推断类型。

修改后的代码 #

from __future__ import annotations
from typing import Generic, TypeVar

# 定义类型变量，绑定到Bar及其子类
T = TypeVar('T', bound='Bar')

class Foo(Generic[T]):
    def __init__(self, value):
        self.value = value

    @property
    def _bar(self) -> type[T]:
        return Bar  # 类型检查会通过，因为Bar是T的边界类型

    def to_bar(self) -> T:
        return self._bar(self.value)

class Bar:
    def __init__(self, value):
        self.value = value

class MyFoo(Foo[MyBar]):
    @property
    def _bar(self) -> type[MyBar]:
        return MyBar

class MyBar(Bar):
    ...

reveal_type(Foo(3).to_bar())
reveal_type(MyFoo(3).to_bar())

mypy检测结果 #

$ mypy t.py
t.py:29: note: Revealed type is "t.Bar"
t.py:30: note: Revealed type is "t.MyBar"
Success: no issues found in 1 source file

关键说明 #

• TypeVar('T', bound='Bar')：限制类型变量T只能是Bar或其子类，保证类型安全。
• class Foo(Generic[T])：将Foo声明为泛型类，T作为类的类型参数。
• 把_bar的返回类型改为type[T]，to_bar的返回类型改为T：让父类的方法返回类型与泛型参数绑定，子类继承时指定具体的MyBar作为类型参数，mypy就能自动推断出子类方法的返回类型。

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