NumPy ndarray dtype 的类型提示？

Question

我想要一个函数来包含 NumPy ndarray 的类型提示及其 dtype 。

例如，对于列表，可以执行以下操作……

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

…为了给出类型提示 bar 必须是 list 由 int 组成。

不幸的是，此语法会为 NumPy 抛出异常 ndarray ：

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

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

是否可以为 dtype np.ndarray --- 的特定类型提示？

原文由 daniel451 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

查看 数据科学类型 包。

 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

原文由 kevin_theinfinityfund 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

您可以查看 nptyping ：

 from nptyping import NDArray, Bool

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

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

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

原文由 R H 发布，翻译遵循 CC BY-SA 4.0 许可协议