为支持np.float32/np.float64数组的函数加类型提示遇mypy报错求解决方案

我想要为一个既接受np.float32数组又接受np.float64数组、且返回相同类型的函数添加类型提示，编写了如下代码：

from typing import overload, Union

import numpy as np
import numpy.typing as npt

NPArray_FLOAT32 = npt.NDArray[np.float32]
NPArray_FLOAT64 = npt.NDArray[np.float64]
NPArray_FLOAT32_64 = Union[NPArray_FLOAT32, NPArray_FLOAT64]

@overload
def foo(xa: NPArray_FLOAT32, xb: NPArray_FLOAT32) -> NPArray_FLOAT32: ...


@overload
def foo(xa: NPArray_FLOAT64, xb: NPArray_FLOAT64) -> NPArray_FLOAT64: ...


def foo(xa: NPArray_FLOAT32_64, xb: NPArray_FLOAT32_64) -> NPArray_FLOAT32_64:
    # ...

运行mypy时出现错误：

mypy [overload-overlap]: Overloaded function signatures 1 and 2 overlap with incompatible return types.

我认为np.float32与np.float64并不存在类型重叠，这似乎是mypy的bug，请问正确的实现方式是什么？

原因分析 #

mypy会将np.float32和np.float64视为与Python原生float兼容的类型，而npt.NDArray在类型检查逻辑中，会判定不同数值类型的数组存在潜在的类型重叠可能（比如隐式类型转换场景），因此认为两个重载签名冲突。

正确实现方式 #

方案1：使用泛型（推荐） #

通过定义受限制的类型变量，让mypy自动推导输入输出的一致类型，代码更简洁且符合类型提示最佳实践：

from typing import Generic, TypeVar

import numpy as np
import numpy.typing as npt

# 定义仅允许np.float32或np.float64的类型变量
FloatType = TypeVar("FloatType", np.float32, np.float64)

def foo(xa: npt.NDArray[FloatType], xb: npt.NDArray[FloatType]) -> npt.NDArray[FloatType]:
    # 函数实现示例
    return xa + xb

方案2：重载+抑制mypy警告 #

如果必须保留重载结构，可以使用mypy_extensions的suppress_message装饰器，明确排除重叠误判（要求mypy版本≥0.971）：

from typing import overload, Union
from mypy_extensions import suppress_message

import numpy as np
import numpy.typing as npt

NPArray_FLOAT32 = npt.NDArray[np.float32]
NPArray_FLOAT64 = npt.NDArray[np.float64]
NPArray_FLOAT32_64 = Union[NPArray_FLOAT32, NPArray_FLOAT64]

@suppress_message("overload-overlap")
@overload
def foo(xa: NPArray_FLOAT32, xb: NPArray_FLOAT32) -> NPArray_FLOAT32: ...

@suppress_message("overload-overlap")
@overload
def foo(xa: NPArray_FLOAT64, xb: NPArray_FLOAT64) -> NPArray_FLOAT64: ...

def foo(xa: NPArray_FLOAT32_64, xb: NPArray_FLOAT32_64) -> NPArray_FLOAT32_64:
    # 函数实现示例
    return xa + xb

方案对比 #

• 泛型方案：无需额外依赖，mypy能完全正确推导类型关系，代码更简洁，是推荐的标准写法。
• 重载+抑制警告：适合必须保留重载结构的场景，但本质是绕过检查，需要依赖第三方扩展。

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