如何编写函数类型提示 兼容numpy且支持Python原生数值类型

触发mypy报错的核心原因有两个：

• 手动定义的Union[int, float, np.floating[Any]]属于自定义联合类型，没有和numpy官方类型存根中声明的合法入参类型做关联，mypy不会自动将其识别为np.mean支持的数值类型
• Python的list是类型不变（invariant）的可变容器，当你显式把返回值列表标注为list[自定义Number类型]时，即使联合类型内的所有单个成员都符合np.mean的入参要求，mypy也会判定整个列表类型不兼容。

不需要做类型强制转换，按以下规则写类型标注即可通过校验：

• 优先复用numpy官方维护的类型别名，不要自己零散拼接数值联合类型
• 必要时用@overload明确入参和返回值的对应关系，避免联合类型的推断歧义
• 不要手动给存储返回值的列表加固定的list类型标注，让mypy自动推导序列类型即可

可直接运行的示例代码 #

from typing import Union, overload
import numpy as np
from numpy.typing import ArrayLike

# 单值入参场景用重载明确类型映射
@overload
def add_one(x: int) -> float | np.floating[Any]: ...
@overload
def add_one(x: float) -> float | np.floating[Any]: ...
@overload
def add_one(x: np.floating[Any]) -> np.floating[Any]: ...
def add_one(x: Union[int, float, np.floating[Any]]) -> Union[float, np.floating[Any]]:
    return x + 1

if __name__ == "__main__":
    # 直接传入原生int组成的列表，不需要强转
    raw_nums = [1, 2, 3, 4]
    # 不要手动给res加list[xxx]类型标注，让mypy自动推导
    res = [add_one(num) for num in raw_nums]
    # 直接传入np.mean，无类型报错
    print(np.mean(res))

额外注意事项 #

• 若使用numpy 1.21及以上版本，numpy.typing.ArrayLike已覆盖Python原生int/float、numpy数值标量、numpy数组、普通数值序列所有合法类型，处理批量入参时可直接用ArrayLike作为入参类型，返回值标注为NDArray[np.floating[Any]]，兼容性更强
• 不要为了通过校验把所有入参强制标注为np.floating[Any]，会丢失对原生数值类型的兼容，属于本末倒置的写法

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