Type hint for NumPy ndarray dtype?
Solution 1
Check out data-science-types package.
pip install data-science-types
MyPy now has access to Numpy, Pandas, and Matplotlib stubs. Allows scenarios like:
# 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")
Use mypy like normal.
$ mypy program.py
Usage with function-parameters
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
Solution 2
You could check out nptyping:
from nptyping import NDArray, Bool
def foo(bar: NDArray[Bool]):
...
Or you could just use strings for type hints:
def foo(bar: 'np.ndarray[np.bool]'):
...
Solution 3
To the best of my knowledge it's not possible yet to specify dtype
in numpy array type hints in function signatures. It is planned to be implemented at some point in the future. See numpy GitHub issue #7370 and numpy-stubs GitHub for more details on the current development status.
Solution 4
One informal solution for type documentation is the following:
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)])))
We've been using this at my company and made a MyPy checker that actually checks that the shapes work out (which we should release at some point).
Only thing is it doesn't make PyCharm happy (ie you still get the nasty warning lines):
daniel451
Updated on October 07, 2021Comments
-
daniel451 over 2 years
I would like a function to include a type hint for NumPy
ndarray
's alongside with itsdtype
.With lists, for example, one could do the following...
def foo(bar: List[int]): ...
...in order to give a type hint that
bar
has to belist
consisting ofint
's.Unfortunately, this syntax throws exceptions for NumPy
ndarray
:def foo(bar: np.ndarray[np.bool]): ... > np.ndarray[np.bool]) (...) TypeError: 'type' object is not subscriptable
Is it possible to give
dtype
-specific type hints fornp.ndarray
?