diff --git a/numpy-stubs/__init__.pyi b/numpy-stubs/__init__.pyi index 973a88d..43f0f71 100644 --- a/numpy-stubs/__init__.pyi +++ b/numpy-stubs/__init__.pyi @@ -3,6 +3,8 @@ import sys import datetime as dt from numpy.core._internal import _ctypes +from numpy.typing import ArrayLike, DtypeLike, _Shape, _ShapeLike + from typing import ( Any, ByteString, @@ -36,69 +38,18 @@ else: from typing import SupportsBytes if sys.version_info >= (3, 8): - from typing import Literal + from typing import Literal, Protocol else: - from typing_extensions import Literal + from typing_extensions import Literal, Protocol # TODO: remove when the full numpy namespace is defined def __getattr__(name: str) -> Any: ... -_Shape = Tuple[int, ...] - -# Anything that can be coerced to a shape tuple -_ShapeLike = Union[int, Sequence[int]] - -_DtypeLikeNested = Any # TODO: wait for support for recursive types - -# Anything that can be coerced into numpy.dtype. -# Reference: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html -_DtypeLike = Union[ - dtype, - # default data type (float64) - None, - # array-scalar types and generic types - type, # TODO: enumerate these when we add type hints for numpy scalars - # TODO: add a protocol for anything with a dtype attribute - # character codes, type strings or comma-separated fields, e.g., 'float64' - str, - # (flexible_dtype, itemsize) - Tuple[_DtypeLikeNested, int], - # (fixed_dtype, shape) - Tuple[_DtypeLikeNested, _ShapeLike], - # [(field_name, field_dtype, field_shape), ...] - # - # The type here is quite broad because NumPy accepts quite a wide - # range of inputs inside the list; see the tests for some - # examples. - List[Any], - # {'names': ..., 'formats': ..., 'offsets': ..., 'titles': ..., - # 'itemsize': ...} - # TODO: use TypedDict when/if it's officially supported - Dict[ - str, - Union[ - Sequence[str], # names - Sequence[_DtypeLikeNested], # formats - Sequence[int], # offsets - Sequence[Union[bytes, Text, None]], # titles - int, # itemsize - ], - ], - # {'field1': ..., 'field2': ..., ...} - Dict[str, Tuple[_DtypeLikeNested, int]], - # (base_dtype, new_dtype) - Tuple[_DtypeLikeNested, _DtypeLikeNested], -] - _NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray) -_ArrayLike = TypeVar("_ArrayLike") - class dtype: names: Optional[Tuple[str, ...]] - def __init__( - self, obj: _DtypeLike, align: bool = ..., copy: bool = ... - ) -> None: ... + def __init__(self, obj: DtypeLike, align: bool = ..., copy: bool = ...) -> None: ... @property def alignment(self) -> int: ... @property @@ -217,6 +168,7 @@ class _ArrayOrScalarCommon( def shape(self) -> _Shape: ... @property def strides(self) -> _Shape: ... + def __array__(self, __dtype: DtypeLike = ...) -> ndarray: ... def __int__(self) -> int: ... def __float__(self) -> float: ... def __complex__(self) -> complex: ... @@ -299,7 +251,7 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container): def __new__( cls, shape: Sequence[int], - dtype: Union[_DtypeLike, str] = ..., + dtype: DtypeLike = ..., buffer: _BufferType = ..., offset: int = ..., strides: _ShapeLike = ..., @@ -338,7 +290,7 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container): def dumps(self) -> bytes: ... def astype( self, - dtype: _DtypeLike, + dtype: DtypeLike, order: str = ..., casting: str = ..., subok: bool = ..., @@ -349,14 +301,14 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container): @overload def view(self, dtype: Type[_NdArraySubClass]) -> _NdArraySubClass: ... @overload - def view(self, dtype: _DtypeLike = ...) -> ndarray: ... + def view(self, dtype: DtypeLike = ...) -> ndarray: ... @overload def view( - self, dtype: _DtypeLike, type: Type[_NdArraySubClass] + self, dtype: DtypeLike, type: Type[_NdArraySubClass] ) -> _NdArraySubClass: ... @overload def view(self, *, type: Type[_NdArraySubClass]) -> _NdArraySubClass: ... - def getfield(self, dtype: Union[_DtypeLike, str], offset: int = ...) -> ndarray: ... + def getfield(self, dtype: DtypeLike, offset: int = ...) -> ndarray: ... def setflags( self, write: bool = ..., align: bool = ..., uic: bool = ... ) -> None: ... @@ -501,26 +453,26 @@ class str_(character): ... def array( object: object, - dtype: _DtypeLike = ..., + dtype: DtypeLike = ..., copy: bool = ..., subok: bool = ..., ndmin: int = ..., ) -> ndarray: ... def zeros( - shape: _ShapeLike, dtype: _DtypeLike = ..., order: Optional[str] = ... + shape: _ShapeLike, dtype: DtypeLike = ..., order: Optional[str] = ... ) -> ndarray: ... def ones( - shape: _ShapeLike, dtype: _DtypeLike = ..., order: Optional[str] = ... + shape: _ShapeLike, dtype: DtypeLike = ..., order: Optional[str] = ... ) -> ndarray: ... def zeros_like( - a: _ArrayLike, + a: ArrayLike, dtype: Optional[dtype] = ..., order: str = ..., subok: bool = ..., shape: Optional[Union[int, Sequence[int]]] = ..., ) -> ndarray: ... def ones_like( - a: _ArrayLike, + a: ArrayLike, dtype: Optional[dtype] = ..., order: str = ..., subok: bool = ..., @@ -530,7 +482,7 @@ def full( shape: _ShapeLike, fill_value: Any, dtype: Optional[dtype] = ..., order: str = ... ) -> ndarray: ... def full_like( - a: _ArrayLike, + a: ArrayLike, fill_value: Any, dtype: Optional[dtype] = ..., order: str = ..., @@ -538,35 +490,35 @@ def full_like( shape: Optional[_ShapeLike] = ..., ) -> ndarray: ... def count_nonzero( - a: _ArrayLike, axis: Optional[Union[int, Tuple[int], Tuple[int, int]]] = ... + a: ArrayLike, axis: Optional[Union[int, Tuple[int], Tuple[int, int]]] = ... ) -> Union[int, ndarray]: ... def isfortran(a: ndarray) -> bool: ... -def argwhere(a: _ArrayLike) -> ndarray: ... -def flatnonzero(a: _ArrayLike) -> ndarray: ... -def correlate(a: _ArrayLike, v: _ArrayLike, mode: str = ...) -> ndarray: ... -def convolve(a: _ArrayLike, v: _ArrayLike, mode: str = ...) -> ndarray: ... -def outer(a: _ArrayLike, b: _ArrayLike, out: ndarray = ...) -> ndarray: ... +def argwhere(a: ArrayLike) -> ndarray: ... +def flatnonzero(a: ArrayLike) -> ndarray: ... +def correlate(a: ArrayLike, v: ArrayLike, mode: str = ...) -> ndarray: ... +def convolve(a: ArrayLike, v: ArrayLike, mode: str = ...) -> ndarray: ... +def outer(a: ArrayLike, b: ArrayLike, out: ndarray = ...) -> ndarray: ... def tensordot( - a: _ArrayLike, - b: _ArrayLike, + a: ArrayLike, + b: ArrayLike, axes: Union[ int, Tuple[int, int], Tuple[Tuple[int, int], ...], Tuple[List[int, int], ...] ] = ..., ) -> ndarray: ... def roll( - a: _ArrayLike, + a: ArrayLike, shift: Union[int, Tuple[int, ...]], axis: Optional[Union[int, Tuple[int, ...]]] = ..., ) -> ndarray: ... -def rollaxis(a: _ArrayLike, axis: int, start: int = ...) -> ndarray: ... +def rollaxis(a: ArrayLike, axis: int, start: int = ...) -> ndarray: ... def moveaxis( a: ndarray, source: Union[int, Sequence[int]], destination: Union[int, Sequence[int]], ) -> ndarray: ... def cross( - a: _ArrayLike, - b: _ArrayLike, + a: ArrayLike, + b: ArrayLike, axisa: int = ..., axisb: int = ..., axisc: int = ..., @@ -581,21 +533,21 @@ def binary_repr(num: int, width: Optional[int] = ...) -> str: ... def base_repr(number: int, base: int = ..., padding: int = ...) -> str: ... def identity(n: int, dtype: Optional[dtype] = ...) -> ndarray: ... def allclose( - a: _ArrayLike, - b: _ArrayLike, + a: ArrayLike, + b: ArrayLike, rtol: float = ..., atol: float = ..., equal_nan: bool = ..., ) -> bool: ... def isclose( - a: _ArrayLike, - b: _ArrayLike, + a: ArrayLike, + b: ArrayLike, rtol: float = ..., atol: float = ..., equal_nan: bool = ..., ) -> Union[bool_, ndarray]: ... -def array_equal(a1: _ArrayLike, a2: _ArrayLike) -> bool: ... -def array_equiv(a1: _ArrayLike, a2: _ArrayLike) -> bool: ... +def array_equal(a1: ArrayLike, a2: ArrayLike) -> bool: ... +def array_equiv(a1: ArrayLike, a2: ArrayLike) -> bool: ... # # Constants @@ -649,7 +601,7 @@ class ufunc: def __name__(self) -> str: ... def __call__( self, - *args: _ArrayLike, + *args: ArrayLike, out: Optional[Union[ndarray, Tuple[ndarray, ...]]] = ..., where: Optional[ndarray] = ..., # The list should be a list of tuples of ints, but since we @@ -664,7 +616,7 @@ class ufunc: casting: str = ..., # TODO: make this precise when we can use Literal. order: Optional[str] = ..., - dtype: Optional[_DtypeLike] = ..., + dtype: DtypeLike = ..., subok: bool = ..., signature: Union[str, Tuple[str]] = ..., # In reality this should be a length of list 3 containing an @@ -876,7 +828,7 @@ def take( ) -> _ScalarNumpy: ... @overload def take( - a: _ArrayLike, + a: ArrayLike, indices: int, axis: Optional[int] = ..., out: Optional[ndarray] = ..., @@ -884,48 +836,48 @@ def take( ) -> _ScalarNumpy: ... @overload def take( - a: _ArrayLike, + a: ArrayLike, indices: _ArrayLikeIntOrBool, axis: Optional[int] = ..., out: Optional[ndarray] = ..., mode: _Mode = ..., ) -> Union[_ScalarNumpy, ndarray]: ... -def reshape(a: _ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ... +def reshape(a: ArrayLike, newshape: _ShapeLike, order: _Order = ...) -> ndarray: ... @overload def choose( a: _ScalarIntOrBool, - choices: Union[Sequence[_ArrayLike], ndarray], + choices: Union[Sequence[ArrayLike], ndarray], out: Optional[ndarray] = ..., mode: _Mode = ..., ) -> _ScalarIntOrBool: ... @overload def choose( a: _IntOrBool, - choices: Union[Sequence[_ArrayLike], ndarray], + choices: Union[Sequence[ArrayLike], ndarray], out: Optional[ndarray] = ..., mode: _Mode = ..., ) -> Union[integer, bool_]: ... @overload def choose( a: _ArrayLikeIntOrBool, - choices: Union[Sequence[_ArrayLike], ndarray], + choices: Union[Sequence[ArrayLike], ndarray], out: Optional[ndarray] = ..., mode: _Mode = ..., ) -> ndarray: ... def repeat( - a: _ArrayLike, repeats: _ArrayLikeIntOrBool, axis: Optional[int] = ... + a: ArrayLike, repeats: _ArrayLikeIntOrBool, axis: Optional[int] = ... ) -> ndarray: ... def put( - a: ndarray, ind: _ArrayLikeIntOrBool, v: _ArrayLike, mode: _Mode = ... + a: ndarray, ind: _ArrayLikeIntOrBool, v: ArrayLike, mode: _Mode = ... ) -> None: ... def swapaxes( - a: Union[Sequence[_ArrayLike], ndarray], axis1: int, axis2: int + a: Union[Sequence[ArrayLike], ndarray], axis1: int, axis2: int ) -> ndarray: ... def transpose( - a: _ArrayLike, axes: Union[None, Sequence[int], ndarray] = ... + a: ArrayLike, axes: Union[None, Sequence[int], ndarray] = ... ) -> ndarray: ... def partition( - a: _ArrayLike, + a: ArrayLike, kth: _ArrayLikeIntOrBool, axis: Optional[int] = ..., kind: _PartitionKind = ..., @@ -949,20 +901,20 @@ def argpartition( ) -> ndarray: ... @overload def argpartition( - a: _ArrayLike, + a: ArrayLike, kth: _ArrayLikeIntOrBool, axis: Optional[int] = ..., kind: _PartitionKind = ..., order: Union[None, str, Sequence[str]] = ..., ) -> ndarray: ... def sort( - a: Union[Sequence[_ArrayLike], ndarray], + a: Union[Sequence[ArrayLike], ndarray], axis: Optional[int] = ..., kind: Optional[_SortKind] = ..., order: Union[None, str, Sequence[str]] = ..., ) -> ndarray: ... def argsort( - a: Union[Sequence[_ArrayLike], ndarray], + a: Union[Sequence[ArrayLike], ndarray], axis: Optional[int] = ..., kind: Optional[_SortKind] = ..., order: Union[None, str, Sequence[str]] = ..., diff --git a/numpy-stubs/typing.pyi b/numpy-stubs/typing.pyi new file mode 100644 index 0000000..f570519 --- /dev/null +++ b/numpy-stubs/typing.pyi @@ -0,0 +1,64 @@ +import sys +from typing import Any, Dict, List, overload, Sequence, Text, Tuple, Union + +from numpy import dtype, ndarray + +if sys.version_info >= (3, 8): + from typing import Protocol +else: + from typing_extensions import Protocol + +_Shape = Tuple[int, ...] + +# Anything that can be coerced to a shape tuple +_ShapeLike = Union[int, Sequence[int]] + +_DtypeLikeNested = Any # TODO: wait for support for recursive types + +# Anything that can be coerced into numpy.dtype. +# Reference: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html +DtypeLike = Union[ + dtype, + # default data type (float64) + None, + # array-scalar types and generic types + type, # TODO: enumerate these when we add type hints for numpy scalars + # TODO: add a protocol for anything with a dtype attribute + # character codes, type strings or comma-separated fields, e.g., 'float64' + str, + # (flexible_dtype, itemsize) + Tuple[_DtypeLikeNested, int], + # (fixed_dtype, shape) + Tuple[_DtypeLikeNested, _ShapeLike], + # [(field_name, field_dtype, field_shape), ...] + # + # The type here is quite broad because NumPy accepts quite a wide + # range of inputs inside the list; see the tests for some + # examples. + List[Any], + # {'names': ..., 'formats': ..., 'offsets': ..., 'titles': ..., + # 'itemsize': ...} + # TODO: use TypedDict when/if it's officially supported + Dict[ + str, + Union[ + Sequence[str], # names + Sequence[_DtypeLikeNested], # formats + Sequence[int], # offsets + Sequence[Union[bytes, Text, None]], # titles + int, # itemsize + ], + ], + # {'field1': ..., 'field2': ..., ...} + Dict[str, Tuple[_DtypeLikeNested, int]], + # (base_dtype, new_dtype) + Tuple[_DtypeLikeNested, _DtypeLikeNested], +] + +class _SupportsArray(Protocol): + @overload + def __array__(self, __dtype: DtypeLike = ...) -> ndarray: ... + @overload + def __array__(self, dtype: DtypeLike = ...) -> ndarray: ... + +ArrayLike = Union[bool, int, float, complex, _SupportsArray, Sequence] diff --git a/tests/fail/array_like.py b/tests/fail/array_like.py new file mode 100644 index 0000000..a5ef579 --- /dev/null +++ b/tests/fail/array_like.py @@ -0,0 +1,22 @@ +from typing import Any, TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from numpy.typing import ArrayLike +else: + ArrayLike = Any + + +class A: + pass + + +x1: ArrayLike = (i for i in range(10)) # E: Incompatible types in assignment +x2: ArrayLike = A() # E: Incompatible types in assignment +x3: ArrayLike = {1: "foo", 2: "bar"} # E: Incompatible types in assignment + +scalar = np.int64(1) +scalar.__array__(dtype=np.float64) # E: Unexpected keyword argument +array = np.array([1]) +array.__array__(dtype=np.float64) # E: Unexpected keyword argument diff --git a/tests/pass/array_like.py b/tests/pass/array_like.py new file mode 100644 index 0000000..098149c --- /dev/null +++ b/tests/pass/array_like.py @@ -0,0 +1,44 @@ +from typing import Any, List, Optional, TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from numpy.typing import ArrayLike, DtypeLike, _SupportsArray +else: + ArrayLike = Any + DtypeLike = Any + _SupportsArray = Any + +x1: ArrayLike = True +x2: ArrayLike = 5 +x3: ArrayLike = 1.0 +x4: ArrayLike = 1 + 1j +x5: ArrayLike = np.int8(1) +x6: ArrayLike = np.float64(1) +x7: ArrayLike = np.complex128(1) +x8: ArrayLike = np.array([1, 2, 3]) +x9: ArrayLike = [1, 2, 3] +x10: ArrayLike = (1, 2, 3) +x11: ArrayLike = "foo" + + +class A: + def __array__(self, dtype: DtypeLike = None) -> np.ndarray: + return np.array([1, 2, 3]) + + +x12: ArrayLike = A() + +scalar: _SupportsArray = np.int64(1) +scalar.__array__(np.float64) +array: _SupportsArray = np.array(1) +array.__array__(np.float64) + +a: _SupportsArray = A() +a.__array__(np.int64) +a.__array__(dtype=np.int64) + +# Escape hatch for when you mean to make something like an object +# array. +object_array_scalar: Any = (i for i in range(10)) +np.array(object_array_scalar) diff --git a/tests/reveal/fromnumeric.py b/tests/reveal/fromnumeric.py index 3894c6b..1fde52c 100644 --- a/tests/reveal/fromnumeric.py +++ b/tests/reveal/fromnumeric.py @@ -23,12 +23,12 @@ np.take(B, 0) # E: Union[numpy.generic, datetime.datetime, datetime.timedelta] ) reveal_type( - np.take( # E: Union[numpy.generic, datetime.datetime, datetime.timedelta, numpy.ndarray] + np.take( # E: Union[Union[numpy.generic, datetime.datetime, datetime.timedelta], numpy.ndarray] A, [0] ) ) reveal_type( - np.take( # E: Union[numpy.generic, datetime.datetime, datetime.timedelta, numpy.ndarray] + np.take( # E: Union[Union[numpy.generic, datetime.datetime, datetime.timedelta], numpy.ndarray] B, [0] ) )