MAINT: add dtype argument to _SupportsArray protocol

Author: person142

numpy-stubs/__init__.pyi

@@ -93,9 +93,21 @@ _DtypeLike = Union[
 _NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray)
 
 class _SupportsArray(Protocol):
-    def __array__(self) -> Union[Sequence, ndarray]: ...
+    def __array__(
+        self, dtype: Optional[_DtypeLike] = ...
+    ) -> Union[Sequence, ndarray]: ...
 
-ArrayLike = Union[bool, int, float, complex, _SupportsArray, Sequence]
+# Note that generic.__array__ and ndarray.__array__ don't accept
+# keyword arguments
+#
+# >>> np.ndarray.__array__(np.array([1]), dtype=np.float64)
+# Traceback (most recent call last):
+#   File "<stdin>", line 1, in <module>
+# TypeError: __array__() takes no keyword arguments
+#
+# so they don't actually satisfy the protocol we've defined here. For
+# that reason include them explicitly in the union.
+ArrayLike = Union[bool, int, float, complex, generic, ndarray, _SupportsArray, Sequence]
 
 class dtype:
     names: Optional[Tuple[str, ...]]
@@ -220,7 +232,6 @@ 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/fail/array_like_py3.py

@@ -1,7 +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
+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

@@ -2,29 +2,33 @@
 
 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'
+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):
+    def __array__(self, dtype=None):
         return [1, 2, 3]
 
-x12: 'np.ArrayLike' = A()
+
+x12: "np.ArrayLike" = A()
+
 
 class B:
-    def __array__(self):
+    def __array__(self, dtype=None):
         return np.array([1, 2, 3])
 
-x13: 'np.ArrayLike' = B()
+
+x13: "np.ArrayLike" = B()
 
 # Escape hatch for when you mean to make something like an object
 # array.