NumPy ndarray dtype 的類型提示？

Question

我想一個功能包括NumPy的一種暗示ndarray “與以S一起dtype 。

例如，對於列表，您可以執行以下操作...

def foo(bar: List[int]):
   ...

...為了給出類型提示， bar必須是由int組成的list 。

不幸的是，這種語法會引發 NumPy ndarray異常：

def foo(bar: np.ndarray[np.bool]):
   ...

> np.ndarray[np.bool]) (...) TypeError: 'type' object is not subscriptable

是否有可能給dtype的特異類型提示np.ndarray ？

Answer 1

您可以查看nptyping ：

from nptyping import NDArray, Bool

def foo(bar: NDArray[Bool]):
   ...

或者你可以只使用字符串作為類型提示：

def foo(bar: 'np.ndarray[np.bool]'):
   ...

Answer 2

查看 數據科學類型包。

pip install data-science-types

MyPy 現在可以訪問 Numpy、Pandas 和 Matplotlib 存根。 允許以下場景：

# program.py

import numpy as np
import pandas as pd

arr1: np.ndarray[np.int64] = np.array([3, 7, 39, -3])  # OK
arr2: np.ndarray[np.int32] = np.array([3, 7, 39, -3])  # Type error

df: pd.DataFrame = pd.DataFrame({'col1': [1,2,3], 'col2': [4,5,6]}) # OK
df1: pd.DataFrame = pd.Series([1,2,3]) # error: Incompatible types in assignment (expression has type "Series[int]", variable has type "DataFrame")

像平常一樣使用 mypy。

$ mypy program.py

與函數參數一起使用

def f(df: pd.DataFrame):
    return df.head()

if __name__ == "__main__":
    x = pd.DataFrame({'col1': [1, 2, 3, 4, 5, 6]})
    print(f(x))

$ mypy program.py
> Success: no issues found in 1 source file

Answer 3

據我所知，目前還不可能在函數簽名的 numpy 數組類型提示中指定dtype 。 計划在未來的某個時間點實施。 有關當前開發狀態​​的更多詳細信息，請參閱numpy GitHub 問題 #7370和numpy-stubs GitHub 。

Answer 4

類型文檔的一種非正式解決方案如下：

from typing import TypeVar, Generic, Tuple, Union, Optional
import numpy as np

Shape = TypeVar("Shape")
DType = TypeVar("DType")


class Array(np.ndarray, Generic[Shape, DType]):
    """
    Use this to type-annotate numpy arrays, e.g.

        def transform_image(image: Array['H,W,3', np.uint8], ...):
            ...

    """
    pass


def func(arr: Array['N,2', int]):
    return arr*2


print(func(arr = np.array([(1, 2), (3, 4)])))

我們一直在我的公司使用它，並制作了一個 MyPy 檢查器來實際檢查形狀是否有效（我們應該在某個時候發布）。

唯一的事情是它不會讓 PyC​​harm 高興（即你仍然會收到令人討厭的警告線）：