如何为NumPy浮点向量/矩阵添加正确的类型注解？

核心方案：结合numpy.typing.NDArray与类型约束 #

要实现指定维度（1维向量、2维矩阵）且支持float16/float32/float64的类型注解，可按以下步骤操作：

1. 导入必要模块

首先引入NumPy及类型注解所需的工具：

import numpy as np
from numpy.typing import NDArray
from typing import Union

2. 定义复用的类型别名

为简化注解代码，先定义支持的浮点类型集合及对应维度的数组类型：

# 涵盖所有目标浮点类型
FloatDType = Union[np.float16, np.float32, np.float64]

# 1维浮点向量（类型注解层面标记为NDArray，运行时需自行校验维度）
FloatVector = NDArray[FloatDType]

# 2维浮点矩阵
FloatMatrix = NDArray[FloatDType]

3. 为函数添加注解

直接使用上面定义的类型别名标注参数：

def process_float_arrays(myVectorFloat: FloatVector, myMatrixFloat: FloatMatrix) -> None:
    # 注意：类型注解仅用于静态检查，不会在运行时强制维度约束
    # 若需严格校验，添加运行时断言
    assert len(myVectorFloat.shape) == 1, "参数必须是1维向量"
    assert len(myMatrixFloat.shape) == 2, "参数必须是2维矩阵"
    
    # 函数业务逻辑示例
    print(f"向量形状：{myVectorFloat.shape}")
    print(f"矩阵形状：{myMatrixFloat.shape}")

进阶：让类型检查工具识别维度

如果希望mypy等静态类型检查工具能直接识别维度约束，可使用Annotated结合Shape（Python 3.11+可直接从typing导入，低版本需安装typing_extensions）：

from typing_extensions import Annotated, Shape  # Python <3.11
# from typing import Annotated, Shape  # Python >=3.11

def process_float_arrays(
    myVectorFloat: Annotated[NDArray[FloatDType], Shape["N"]],
    myMatrixFloat: Annotated[NDArray[FloatDType], Shape["N", "M"]]
) -> None:
    # 运行时维度校验仍需保留
    assert len(myVectorFloat.shape) == 1
    assert len(myMatrixFloat.shape) == 2
    # ...

关键说明 #

• NDArray[DType]是NumPy官方推荐的数组类型注解方式，替代旧版的np.ndarray
• 类型注解仅用于静态代码检查，不会自动在运行时校验维度或类型，需手动添加assert或自定义校验逻辑
• Union[np.float16, np.float32, np.float64]明确指定了支持的浮点类型，若需兼容所有浮点子类，也可使用np.floating[np.number]替代

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