如何解决Python高阶函数返回overload函数时类型提示失效问题？

问题根源在于mypy无法自动推断高阶函数内部定义的重载函数的类型签名，需要显式为返回的重载函数定义类型协议，再标注高阶函数的返回类型。

具体解决方案 #

1. 导入Protocol，定义包含重载调用签名的协议类，描述返回函数的所有重载行为；
2. 为高阶函数get_f添加返回类型注解，指定为上述协议类型。

修改后的代码如下：

from typing import Optional, overload, Protocol

class FType(Protocol):
    @overload
    def __call__(self, x: int) -> str: ...

    @overload
    def __call__(self, x: int, c: bool) -> int: ...

def get_f(b: int) -> FType:
    @overload
    def f(x: int) -> str: ...

    @overload
    def f(x: int, c: bool) -> int: ...

    def f(x: int, c: Optional[bool] = None) -> int | str:
        return 2 * x * b if c is not None else "s"

    return f


d = get_f(2)
a = d(2)  # mypy 会推断 a 的类型为 str
b = a ** 2  # mypy 将报错：Unsupported operand types for ** ("str" and "int")

原理说明 #

通过Protocol协议类明确描述返回函数的重载调用签名后，mypy可以依据这个协议解析返回函数的类型，从而正确检测到类型错误。

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