鍵入包含 unsignedinteger 的 numpy.ndarray 的提示

Question

我有一個包含 unsignedinteger 的 numpy 數組，如下所示：

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

array([   5,  100, 2000], dtype=uint16)

此arr將輸入到 function。我想知道 function 參數的類型提示應該是什么樣子？

def myfunc(arr: ?):
    pass

我首先認為它應該是arr: np.ndarray 。 但隨后mypy抱怨了。
Argument "arr" to "myfunc" has incompatible type "unsignedinteger[_16Bit]"; expected "ndarray[Any, Any]" [arg-type]

arr: np.ndarray[np.uint16]也不起作用。
error: "ndarray" expects 2 type arguments, but 1 given [type-arg]

Answer 1

您可以使用numpy中的typing模塊：

import numpy as np
import numpy.typing as npt

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

Answer 2

NumPy 文檔已涵蓋。 您正在尋找numpy.typing.NDArray ：

from numpy import typing as npt

def myfunc(arr: npt.NDArray):
    pass