如何为含无符号整数的numpy.ndarray编写函数类型提示？

针对你遇到的问题，以下是正确的类型提示写法：

推荐写法：使用numpy.typing.NDArray #

这是NumPy官方推荐的简洁写法，专门用于标注数组类型：

import numpy as np
from numpy.typing import NDArray

def myfunc(arr: NDArray[np.uint16]):
    pass

# 测试调用
arr = np.uint16([5, 100, 2000])
myfunc(arr)

完整写法：显式指定np.ndarray的两个类型参数 #

np.ndarray需要两个类型参数：第一个表示数组的形状（用Any表示任意形状），第二个表示元素的dtype类型：

import numpy as np
from typing import Any

def myfunc(arr: np.ndarray[Any, np.dtype[np.uint16]]):
    pass

arr = np.uint16([5, 100, 2000])
myfunc(arr)

为什么之前的写法报错？ #

• 仅用np.ndarray：mypy无法识别数组的具体dtype，若代码中变量的类型推断与通用ndarray不匹配（比如误传了标量而非数组），就会触发类型不兼容错误。
• 用np.ndarray[np.uint16]：np.ndarray要求必须传入两个类型参数（形状和dtype），只传一个不符合类型定义规则，因此报错。

注意：确保你的环境中安装了NumPy的类型支持，通常安装NumPy时会自带，若mypy仍有问题，可以安装types-numpy包补充类型 stub。

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