ENH: get more specific about _ArrayLike, make it public

Closes #37.

Add tests to check various examples. Note that supporting __array__
also requires making _DtypeLike public too, so this does that as well.

Author: person142

numpy-stubs/__init__.pyi

@@ -36,9 +36,9 @@ 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: ...
@@ -52,7 +52,7 @@ _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[
+DtypeLike = Union[
     dtype,
     # default data type (float64)
     None,
@@ -92,13 +92,17 @@ _DtypeLike = Union[
 
 _NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray)
 
-_ArrayLike = TypeVar("_ArrayLike")
+class _SupportsArray(Protocol):
+    @overload
+    def __array__(self, __dtype: DtypeLike = ...) -> ndarray: ...
+    @overload
+    def __array__(self, dtype: Optional[DtypeLike] = ...) -> ndarray: ...
+
+ArrayLike = Union[bool, int, float, complex, _SupportsArray, Sequence]
 
 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 +221,7 @@ class _ArrayOrScalarCommon(
     def shape(self) -> _Shape: ...
     @property
     def strides(self) -> _Shape: ...
+    def __array__(self, __dtype: Optional[DtypeLike] = ...) -> ndarray: ...
     def __int__(self) -> int: ...
     def __float__(self) -> float: ...
     def __complex__(self) -> complex: ...
@@ -299,7 +304,7 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container):
     def __new__(
         cls,
         shape: Sequence[int],
-        dtype: Union[_DtypeLike, str] = ...,
+        dtype: Union[DtypeLike, str] = ...,
         buffer: _BufferType = ...,
         offset: int = ...,
         strides: _ShapeLike = ...,
@@ -338,7 +343,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 +354,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: Union[DtypeLike, str], offset: int = ...) -> ndarray: ...
     def setflags(
         self, write: bool = ..., align: bool = ..., uic: bool = ...
     ) -> None: ...
@@ -484,26 +489,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 = ...,
@@ -513,43 +518,43 @@ 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 = ...,
     subok: bool = ...,
     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 = ...,
@@ -564,21 +569,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
@@ -632,7 +637,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
@@ -647,7 +652,7 @@ class ufunc:
         casting: str = ...,
         # TODO: make this precise when we can use Literal.
         order: Optional[str] = ...,
-        dtype: Optional[_DtypeLike] = ...,
+        dtype: Optional[DtypeLike] = ...,
         subok: bool = ...,
         signature: Union[str, Tuple[str]] = ...,
         # In reality this should be a length of list 3 containing an
@@ -845,74 +850,74 @@ def take(
 ) -> _ScalarNumpy: ...
 @overload
 def take(
-    a: _ArrayLike,
+    a: ArrayLike,
     indices: int,
     axis: Optional[int] = ...,
     out: Optional[ndarray] = ...,
     mode: _Mode = ...,
 ) -> _ScalarNumpy: ...
 @overload
 def take(
-    a: _ArrayLike,
+    a: ArrayLike,
     indices: _ArrayLikeInt,
     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: _ScalarGeneric,
-    choices: Union[Sequence[_ArrayLike], ndarray],
+    choices: Union[Sequence[ArrayLike], ndarray],
     out: Optional[ndarray] = ...,
     mode: _Mode = ...,
 ) -> _ScalarGeneric: ...
 @overload
 def choose(
     a: _Scalar,
-    choices: Union[Sequence[_ArrayLike], ndarray],
+    choices: Union[Sequence[ArrayLike], ndarray],
     out: Optional[ndarray] = ...,
     mode: _Mode = ...,
 ) -> _ScalarNumpy: ...
 @overload
 def choose(
-    a: _ArrayLike,
-    choices: Union[Sequence[_ArrayLike], ndarray],
+    a: ArrayLike,
+    choices: Union[Sequence[ArrayLike], ndarray],
     out: Optional[ndarray] = ...,
     mode: _Mode = ...,
 ) -> ndarray: ...
 def repeat(
-    a: _ArrayLike, repeats: _ArrayLikeInt, axis: Optional[int] = ...
+    a: ArrayLike, repeats: _ArrayLikeInt, axis: Optional[int] = ...
 ) -> ndarray: ...
-def put(a: ndarray, ind: _ArrayLikeInt, v: _ArrayLike, mode: _Mode = ...) -> None: ...
+def put(a: ndarray, ind: _ArrayLikeInt, 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: _ArrayLikeInt,
     axis: Optional[int] = ...,
     kind: _PartitionKind = ...,
     order: Union[None, str, Sequence[str]] = ...,
 ) -> ndarray: ...
 def argpartition(
-    a: _ArrayLike,
+    a: ArrayLike,
     kth: _ArrayLikeInt,
     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]] = ...,

scripts/array_protocol.py

@@ -0,0 +1,3 @@
+from typing_extensions import Protocol
+
+import numpy as np

tests/fail/array_like.py

@@ -0,0 +1,10 @@
+import numpy as np
+
+
+class A:
+    pass
+
+
+x1: "np.ArrayLike" = (i for i in range(10))  # E: Incompatible types in assignment
+x2: "np.ArrayLike" = A()  # E: Incompatible types in assignment
+x3: "np.ArrayLike" = {1: "foo", 2: "bar"}  # E: Incompatible types in assignment

tests/pass/array_like.py

@@ -0,0 +1,30 @@
+from typing import Any, List, Optional
+
+import numpy as np
+
+x1: "np.ArrayLike" = True
+x2: "np.ArrayLike" = 5
+x3: "np.ArrayLike" = 1.0
+x4: "np.ArrayLike" = 1 + 1j
+x5: "np.ArrayLike" = np.int8(1)
+x6: "np.ArrayLike" = np.float64(1)
+x7: "np.ArrayLike" = np.complex128(1)
+x8: "np.ArrayLike" = np.array([1, 2, 3])
+x9: "np.ArrayLike" = [1, 2, 3]
+x10: "np.ArrayLike" = (1, 2, 3)
+x11: "np.ArrayLike" = "foo"
+
+
+class B:
+    def __array__(self, dtype: Optional["np.DtypeLike"] = None) -> np.ndarray:
+        return np.array([1, 2, 3])
+
+
+x12: "np.ArrayLike" = B()
+x13: "np._SupportsArray" = np.int64(1)
+x14: "np._SupportsArray" = np.array(1)
+
+# 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)

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]
     )
 )