MAINT: use a _SupportsArray protocol to improve ArrayLike

Author: person142

numpy-stubs/__init__.pyi

@@ -35,6 +35,11 @@ if sys.version_info[0] < 3:
 else:
     from typing import SupportsBytes
 
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
 # TODO: remove when the full numpy namespace is defined
 def __getattr__(name: str) -> Any: ...
 
@@ -87,7 +92,10 @@ _DtypeLike = Union[
 
 _NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray)
 
-ArrayLike = Union[int, float, complex, generic, ndarray, Sequence]
+class _SupportsArray(Protocol):
+    def __array__(self) -> Union[Sequence, ndarray]: ...
+
+ArrayLike = Union[bool, int, float, complex, _SupportsArray, Sequence]
 
 class dtype:
     names: Optional[Tuple[str, ...]]
@@ -212,6 +220,7 @@ class _ArrayOrScalarCommon(
     def shape(self) -> _Shape: ...
     @property
     def strides(self) -> _Shape: ...
+    def __array__(self) -> ndarray: ...
     def __int__(self) -> int: ...
     def __float__(self) -> float: ...
     def __complex__(self) -> complex: ...

tests/pass/array_like_py3.py

@@ -1,12 +1,32 @@
+from typing import Any
+
 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'
+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 A:
+    def __array__(self):
+        return [1, 2, 3]
+
+x12: 'np.ArrayLike' = A()
+
+class B:
+    def __array__(self):
+        return np.array([1, 2, 3])
+
+x13: 'np.ArrayLike' = B()
+
+# 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)