ENH: get more specific about _ArrayLike, make it public

Closes #37.

Add tests to check various examples.

Author: person142

numpy-stubs/__init__.pyi

@@ -87,7 +87,7 @@ _DtypeLike = Union[
 
 _NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray)
 
-_ArrayLike = TypeVar("_ArrayLike")
+ArrayLike = Union[int, float, complex, generic, ndarray, Sequence]
 
 class dtype:
     names: Optional[Tuple[str, ...]]
@@ -491,14 +491,14 @@ def ones(
     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 = ...,
@@ -508,43 +508,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 = ...,
@@ -559,21 +559,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
@@ -627,7 +627,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

tests/fail/array_like_py3.py

@@ -0,0 +1,7 @@
+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_py3.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+x1: 'np.ArrayLike' = 5
+x2: 'np.ArrayLike' = 1.0
+x3: 'np.ArrayLike' = 1 + 1j
+x4: 'np.ArrayLike' = np.int8(1)
+x5: 'np.ArrayLike' = np.float64(1)
+x6: 'np.ArrayLike' = np.complex128(1)
+x7: 'np.ArrayLike' = np.array([1, 2, 3])
+x8: 'np.ArrayLike' = [1, 2, 3]
+x9: 'np.ArrayLike' = (1, 2, 3)
+x10: 'np.ArrayLike' = 'foo'