如何为通用Numpy数组添加类型提示？解决泛型类型报错问题

问题分析 #

你遇到的错误根源在于TypeVar未被正确约束到Numpy的标量类型体系中，类型检查器无法识别E是NDArray接受的合法类型参数。直接绑定具体dtype会丢失泛型灵活性，同时无法覆盖所有支持乘法操作的Numpy类型。

解决方案 #

通过给TypeVar添加bound约束，限定它只能是Numpy的标量类型（npt.ScalarType），既保留泛型能力，又能让类型检查器正确识别类型参数的合法性，同时兼容乘法操作的类型推导：

import numpy as np
import numpy.typing as npt
from typing import TypeVar

# 约束E为Numpy标量类型的子类
E = TypeVar("E", bound=npt.ScalarType)

def double_arr(arr: npt.NDArray[E]) -> npt.NDArray[E]:
    return arr * 2

验证效果 #

修改后类型检查器可以正确推断输入输出的数组类型：

# 输入int8数组，返回类型为npt.NDArray[np.int8]
arr_int = np.array([1, 2, 3], dtype=np.int8)
result_int = double_arr(arr_int)

# 输入float32数组，返回类型为npt.NDArray[np.float32]
arr_float = np.array([1, 2.3, 3], dtype=np.float32)
result_float = double_arr(arr_float)

补充说明 #

• npt.ScalarType是Numpy类型系统中所有标量类型的父类型，涵盖np.int8、np.float32等所有基础dtype对应的标量类型，确保类型检查器能识别所有合法的Numpy数组元素类型。
• 该方案完全兼容Numpy 1.23.5和Python 3.10环境，既满足泛型需求，又能正确处理乘法操作的类型兼容性。

内容的提问来源于stack exchange，提问作者Jorge Ruiz Gómez