如何基于参数推断Python对象的函数返回类型？

假设有如下函数：

def get_attr_wrapper(obj: object, attr: str) -> ???:
    return getattr(obj, attr)

请问如何根据传入的参数推断get_attr_wrapper的返回类型？是否可以通过泛型实现？

例如，当传入以下参数时：

from dataclasses import dataclass

@dataclass
class Foo:
    bar: str

foo = Foo(bar="baz")

rv = get_attr_wrapper(foo, "bar")

期望Python类型检查器能将rv推断为str类型。

完全可以通过泛型实现类型推断，核心是把函数的obj参数从固定的object类型改为泛型类型，让类型检查器能基于传入对象的具体类型推导属性类型。

具体实现 #

修改函数定义，引入TypeVar来表示传入对象的类型：

from typing import TypeVar

T = TypeVar('T')

def get_attr_wrapper(obj: T, attr: str) -> object:
    return getattr(obj, attr)

如果要更清晰地关联返回类型，也可以用两个泛型变量分别表示对象类型和返回类型：

from typing import TypeVar

T = TypeVar('T')
R = TypeVar('R')

def get_attr_wrapper(obj: T, attr: str) -> R:
    return getattr(obj, attr)

类型推断效果 #

当你传入示例中的foo（Foo类型）和"bar"时，mypy、pyright这类类型检查器会识别到Foo类有bar: str的属性，自动将rv的类型推断为str。

关键原理 #

原函数中obj被标注为object类型，类型检查器无法获取该对象的具体属性信息；改用泛型T后，检查器会保留传入对象的原始类型信息，结合属性名就能准确推导返回值的类型。

如果需要更严格的约束（比如限制只能访问对象存在的属性），可以结合Protocol定义具有特定属性的类型，但上述泛型写法已经能满足绝大多数通用场景的类型推断需求。

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