Python 3.12+与NumPy 2+泛型类型提示实现及优化问询

基础类型提示 #

对于无需关注形状的场景，直接使用numpy.ndarray或其别名numpy.typing.NDArray，搭配具体dtype即可：

import numpy as np
from numpy.typing import NDArray

def func(arr: NDArray[np.float64]) -> NDArray[np.float64]:
    return arr + 1.0

带形状的泛型类型提示（NumPy 2+） #

NumPy 2+正式支持PEP 484兼容的泛型，可通过ndarray[DType, Shape]同时指定dtype和形状：

• 固定形状：使用np.Shape["(m, n)"]字符串形式，或typing.Tuple定义维度

def process_2d(arr: np.ndarray[np.int32, np.Shape["(3, 3)"]]) -> np.ndarray[np.int32, np.Shape["(3, 3)"]]:
    return arr @ arr

• 可变形状：借助typing.TypeVarTuple和Unpack定义通用形状参数，适配任意维度的数组

from typing import TypeVarTuple, Unpack
import numpy as np

Shape = TypeVarTuple("Shape")
DType = TypeVar("DType", bound=np.generic)

def scale_array(arr: np.ndarray[DType, Unpack[Shape]], factor: float) -> np.ndarray[DType, Unpack[Shape]]:
    return arr * factor

完全正确。NumPy 2+重构了类型系统，实现了符合PEP 484的泛型支持；而Python 3.12+对泛型语法的支持更完善（无需from __future__ import annotations即可使用字符串形状）。只要确保导入正确的类型（如numpy.typing.NDArray），并使用最新版的mypy和NumPy，泛型写法就能被正常识别。

要让mypy正确解析NumPy 2+的形状类型提示，需注意以下几点：

1. 启用NumPy插件：在mypy.ini或pyproject.toml中配置：

[mypy]
plugins = numpy

2. 形状语法选择：
   • 固定形状优先用np.Shape["(m, n)"]字符串形式，mypy最新版已支持解析
   • 动态形状用Tuple[int, int]（指定维度数）或Tuple[int, ...]（任意维度）
3. 使用NumPy原生dtype：必须用np.float64这类NumPy自带类型，不能混用Python原生的float

不要依赖# type: ignore，可以通过以下方式解决类型检查错误：

1. 显式类型转换：用numpy.typing.cast明确指定数组的类型和形状：

def create_array() -> np.ndarray[np.float64, np.Shape["(2, 2)"]]:
    raw_arr = np.array([[1.0, 2.0], [3.0, 4.0]])
    return np.typing.cast(np.ndarray[np.float64, np.Shape["(2, 2)"]], raw_arr)

2. 更新依赖版本：确保mypy（>=1.5）、NumPy（>=2.0）为最新版，旧版本可能对新语法支持不足
3. 简化类型提示：如果形状检查不是核心需求，暂时用NDArray[DType]忽略形状，只保证dtype正确
4. 微调mypy配置：在配置文件中针对特定规则放宽检查（如disable_error_code = attr-defined），但避免过度关闭严格检查

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