如何为接受实例或构建器并返回实例的函数添加正确类型提示？

原代码的核心问题是两个独立的TypeVar未建立关联，类型检查器无法将FooBuilder.build()的返回类型与函数声明的_F绑定；同时，直接使用bound=Foo的_F时，类型检查器会认为传入的Foo实例可能是_F的基类而非具体子类，因此触发报错。

修正方案 #

我们需要让FooBuilder成为泛型类，明确它构建的Foo子类类型，同时让函数的输入类型与返回类型通过同一个TypeVar关联。

修正后的代码 #

# python 3.9.5
from typing import TypeVar, Union, Generic


class Foo:
    pass


_F = TypeVar("_F", bound=Foo)

class FooBuilder(Generic[_F]):
    def build(self) -> _F:
        return Foo()  # 子类Builder需重写该方法返回对应子类实例


def get_foo(x: Union[_F, FooBuilder[_F]]) -> _F:
    if isinstance(x, Foo):
        return x
    elif isinstance(x, FooBuilder):
        return x.build()


# 基础类型验证
assert isinstance(get_foo(Foo()), Foo)
assert isinstance(get_foo(FooBuilder()), Foo)

# 子类场景验证
class FooSub(Foo):
    pass

class FooSubBuilder(FooBuilder[FooSub]):
    def build(self) -> FooSub:
        return FooSub()

assert isinstance(get_foo(FooSub()), FooSub)
assert isinstance(get_foo(FooSubBuilder()), FooSub)

关键改动说明 #

1. 泛型化FooBuilder：

   • 让FooBuilder继承Generic[_F]，用_F标记它构建的目标类型，build方法明确返回_F。
   • 子类FooSubBuilder可以指定_F为FooSub，并重写build返回对应子类实例，类型检查器能识别这种绑定关系。

2. 调整函数参数类型：

   • 将参数类型改为Union[_F, FooBuilder[_F]]，这样无论传入Foo子类实例，还是能构建该子类的FooBuilder实例，返回类型都统一为_F，类型检查器可以正确推断。

3. 解决类型不匹配问题：

   • 当传入Foo实例时，类型检查器能确认它就是_F（因为_F bound到Foo），返回时不会报错。
   • 当传入FooBuilder[_F]实例时，build()返回的是_F，与函数返回类型完全匹配，消除pylance和mypy的报错。

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