我希望函数包含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
特定类型的提示?
答案 0 :(得分:6)
签出data-science-types软件包。
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
答案 1 :(得分:2)
据我所知,尚无法在函数签名中的numpy数组类型提示中指定dtype
。计划在将来的某个时候实施。有关当前开发状态的更多详细信息,请参见numpy GitHub issue #7370和numpy-stubs GitHub。
答案 2 :(得分:2)
一种非正式的类型文档解决方案如下:
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检查器,该检查器实际上检查形状是否可行(我们应该在某个时候发布)。
唯一的事情是它不能使PyCharm开心(即,您仍然会收到讨厌的警告线):
答案 3 :(得分:1)
您可以签出nptyping:
from nptyping import Array
def foo(bar: Array[np.bool]):
...
或者您可以只使用字符串作为类型提示:
def foo(bar: 'np.ndarray[np.bool]'):
...