如何为numpy类型提示npt.NDArray指定数组轴数（维度）

原生npt.NDArray不支持直接传入轴数作为泛型参数。
npt.NDArray的泛型设计仅接收1个参数，用于指定数组的dtype类型，npt.NDArray[3, np.float64]这类写法不符合numpy官方的类型定义，会直接被mypy、pyright等类型检查器判定为语法错误。

要实现3D float64数组的类型标注，根据对类型校验严格度的要求，有三种常用落地写法：

• 基础写法（无额外依赖，全版本兼容）
  直接使用npt.NDArray指定dtype，通过语义化类型别名明确轴数，搭配运行时维度校验，这是工业界项目最常用的方案：

  import numpy as np
  import numpy.typing as npt
  
  # 定义类型别名，语义明确为3维float64数组
  Float3DArray = npt.NDArray[np.float64]
  
  def process_voxel_data(voxels: Float3DArray) -> None:
      # 运行时校验维度，拦截非法输入
      if voxels.ndim != 3:
          raise ValueError(f"Input must be 3D array, got ndim={voxels.ndim}")
      # 业务逻辑
  

• 静态校验写法（支持pyright、开启numpy插件的mypy）
  借助typing.Annotated附加形状信息，可以在静态检查阶段直接识别轴数，不需要运行时断言就能拦截维度错误的传参：

  from typing import Annotated
  import numpy as np
  import numpy.typing as npt
  
  # 标注为3维float64数组，三个`::`代表不限制每个轴的具体长度，仅校验轴数为3
  Float3DArray = Annotated[npt.NDArray[np.float64], np.ndarray.shape[::, ::, ::]]
  
  def calc_feature_map(features: Float3DArray) -> Float3DArray:
      return features * 2
  

• 维度泛型写法（需要第三方依赖支持）
  如果需要更灵活的维度泛型支持，可以使用第三方形状校验库提供的泛型数组类型，支持任意固定轴数、甚至固定轴长度的数组标注，缺点是需要引入额外依赖。

不要尝试给npt.NDArray传入多个泛型参数，截止numpy 2.0版本，官方类型定义中NDArray仅支持dtype单个泛型入参，多传参会直接报类型错误。

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