如何让mypy从被调用方法推断调用方法的类型签名？

你需要实现Test2.f作为Test2.get_arg().f的快捷方式，同时复用Test1.f的重载类型签名，避免重复编写，且保证mypy能正确推断类型。下面提供两种可行的实现方案：

方案1：利用Protocol定义共享接口 #

通过定义包含f方法重载签名的Protocol，让Test1和Test2都实现该协议，只需写一次重载规则即可实现类型复用。

from typing import overload, Union, Protocol

# 定义包含f方法重载逻辑的协议
class HasF(Protocol):
    @overload
    def f(self, x: int) -> int: ...
    @overload
    def f(self, x: str) -> str: ...

class Test1(HasF):
    # 只需实现一次具体逻辑，自动匹配协议的重载规则
    def f(self, x: Union[int, str]) -> Union[int, str]:
        return x

class Test2(HasF):
    def get_arg(self) -> Test1:
        return Test1()
    
    # 代理调用Test1的f方法，无需重复编写重载签名
    def f(self, x: Union[int, str]) -> Union[int, str]:
        return self.get_arg().f(x)

说明 #

• HasF协议统一定义了f方法的所有重载规则，Test1和Test2通过实现该协议保证类型一致性。
• mypy会自动校验Test2.f的实现是否符合协议要求，由于Test1.f已经满足规则，代理调用自然保证类型安全。
• 仅需在协议中编写一次重载，彻底避免重复代码。

方案2：泛型代理装饰器 #

如果不想显式实现Protocol，可以用泛型装饰器动态复用Test1.f的类型签名，自动适配Test2的self参数。

from typing import overload, Union, ParamSpec, TypeVar, Callable, Protocol

# 定义泛型参数，用于捕获方法的参数和返回值类型
P = ParamSpec("P")
R = TypeVar("R")
T = TypeVar("T")
U = TypeVar("U", bound="HasGetArg[T]")

# 定义协议：要求类必须有get_arg方法返回指定类型
class HasGetArg(Protocol[T]):
    def get_arg(self) -> T: ...

# 泛型代理装饰器：将T类型的方法适配为U类型的方法
def proxy_method(method: Callable[[T, *P.args], R]) -> Callable[[U, *P.args], R]:
    def wrapper(self: U, *args: P.args, **kwargs: P.kwargs) -> R:
        return method(self.get_arg(), *args, **kwargs)
    return wrapper

class Test1:
    @overload
    def f(self, x: int) -> int: ...
    @overload
    def f(self, x: str) -> str: ...
    def f(self, x: Union[int, str]) -> Union[int, str]:
        return x

class Test2(HasGetArg[Test1]):
    def get_arg(self) -> Test1:
        return Test1()
    
    # 直接复用Test1.f的类型签名，无需编写任何重载
    f = proxy_method(Test1.f)

说明 #

• proxy_method装饰器接收Test1.f作为参数，自动生成适配Test2的方法，自动处理self参数的类型差异。
• HasGetArg协议确保Test2具备返回Test1实例的get_arg方法，保证代理调用的类型安全。
• mypy会自动继承Test1.f的所有重载规则，正确推断Test2.f的参数和返回类型，完全无需手动编写类型注解。

验证类型推断 #

用以下代码测试，mypy会正确识别类型：

t2 = Test2()
reveal_type(t2.f(1))  # mypy输出：Revealed type is "builtins.int"
reveal_type(t2.f("hello"))  # mypy输出：Revealed type is "builtins.str"

内容的提问来源于stack exchange，提问作者Zoran Bošnjak