Define special methods for ndarray and add more extensive tests.

numpy/__init__.pyi

@@ -1,15 +1,41 @@
-# very simple, just enough to start running tests
-#
 import builtins
-from typing import Any, Mapping, List, Optional, Tuple, Union
+
+from typing import (
+    Any, Dict, Iterable, List, Optional, Mapping, Sized,
+    SupportsInt, SupportsFloat, SupportsComplex, SupportsBytes, SupportsAbs,
+    Tuple, Union,
+)
+
+import sys
 
 from numpy.core._internal import _ctypes
 
 _Shape = Tuple[int, ...]
 
+# Anything that can be coerced into numpy.dtype. To avoid recursive
+# definitions, any nested fields are required to be castable to a dtype object
+# are typed as Any.
+# Refernce: https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
+_ConvertibleToDtype = Union[
+    type,  # TODO: enumerate np.generic types and Python scalars
+    # TODO: add a protocol for anything with a dtype attribute
+    str,
+    Tuple[Any, int],
+    Tuple[Any, _Shape],
+    List[Union[Tuple[Union[str, Tuple[str, str]], Any],
+               Tuple[Union[str, Tuple[str, str]], Any, _Shape]]],
+    Dict[str, Any],
+    Tuple[Any, Any]]
+
+
 class dtype:
     names: Optional[Tuple[str, ...]]
 
+    def __init__(self,
+                 obj: Union[dtype, _ConvertibleToDtype],
+                 align: bool = ...,
+                 copy: bool = ...) -> None: ...
+
     @property
     def alignment(self) -> int: ...
 
@@ -83,7 +109,8 @@ class dtype:
     def type(self) -> builtins.type: ...
 
 
-_dtype_class = dtype      # for ndarray type
+_DtypeLike = Union[dtype, _ConvertibleToDtype]
+_Dtype = dtype  # to avoid name conflicts with ndarray.dtype
 
 
 class _flagsobj:
@@ -144,8 +171,10 @@ class flatiter:
     def __next__(self) -> Any: ...
 
 
-class ndarray:
-    dtype: _dtype_class
+class ndarray(Iterable, Sized, SupportsInt, SupportsFloat, SupportsComplex,
+              SupportsBytes, SupportsAbs[Any]):
+
+    dtype: _Dtype
     imag: ndarray
     real: ndarray
     shape: _Shape
@@ -181,12 +210,120 @@ class ndarray:
     @property
     def ndim(self) -> int: ...
 
+    # Many of these special methods are irrelevant currently, since protocols
+    # aren't supported yet. That said, I'm adding them for completeness.
+    # https://docs.python.org/3/reference/datamodel.html
+    def __len__(self) -> int: ...
+    def __getitem__(self, key) -> Any: ...
+    def __setitem__(self, key, value): ...
+    def __iter__(self) -> Any: ...
+    def __contains__(self, key) -> bool: ...
+
+    def __int__(self) -> int: ...
+    def __float__(self) -> float: ...
+    def __complex__(self) -> complex: ...
+    if sys.version_info.major < 3:
+        def __oct__(self) -> str: ...
+        def __hex__(self) -> str: ...
+    def __nonzero__(self) -> bool: ...
+    def __bytes__(self) -> bytes: ...
+
+    def __index__(self) -> int: ...
+
+    def __copy__(self, order: str = ...) -> ndarray: ...
+    def __deepcopy__(self, memo: dict) -> ndarray: ...
+
+    # https://github.com/numpy/numpy/blob/v1.13.0/numpy/lib/mixins.py#L63-L181
+
+    # TODO(shoyer): add overloads (returning ndarray) for cases where other is
+    # known not to define __array_priority__ or __array_ufunc__, such as for
+    # numbers or other numpy arrays. Or even better, use protocols (once they
+    # work).
+
+    def __lt__(self, other): ...
+    def __le__(self, other): ...
+    def __eq__(self, other): ...
+    def __ne__(self, other): ...
+    def __gt__(self, other): ...
+    def __ge__(self, other): ...
+
+    def __add__(self, other): ...
+    def __radd__(self, other): ...
+    def __iadd__(self, other): ...
+
+    def __sub__(self, other): ...
+    def __rsub__(self, other): ...
+    def __isub__(self, other): ...
+
+    def __mul__(self, other): ...
+    def __rmul__(self, other): ...
+    def __imul__(self, other): ...
+
+    if sys.version_info.major < 3:
+        def __div__(self, other): ...
+        def __rdiv__(self, other): ...
+        def __idiv__(self, other): ...
+
+    def __truediv__(self, other): ...
+    def __rtruediv__(self, other): ...
+    def __itruediv__(self, other): ...
+
+    def __floordiv__(self, other): ...
+    def __rfloordiv__(self, other): ...
+    def __ifloordiv__(self, other): ...
+
+    def __mod__(self, other): ...
+    def __rmod__(self, other): ...
+    def __imod__(self, other): ...
+
+    def __divmod__(self, other): ...
+    def __rdivmod__(self, other): ...
+
+    # NumPy's __pow__ doesn't handle a third argument
+    def __pow__(self, other): ...
+    def __rpow__(self, other): ...
+    def __ipow__(self, other): ...
+
+    def __lshift__(self, other): ...
+    def __rlshift__(self, other): ...
+    def __ilshift__(self, other): ...
+
+    def __rshift__(self, other): ...
+    def __rrshift__(self, other): ...
+    def __irshift__(self, other): ...
+
+    def __and__(self, other): ...
+    def __rand__(self, other): ...
+    def __iand__(self, other): ...
+
+    def __xor__(self, other): ...
+    def __rxor__(self, other): ...
+    def __ixor__(self, other): ...
+
+    def __or__(self, other): ...
+    def __ror__(self, other): ...
+    def __ior__(self, other): ...
+
+    if sys.version_info >= (3, 5):
+        def __matmul__(self, other): ...
+        def __rmatmul__(self, other): ...
+
+    def __neg__(self) -> ndarray: ...
+    def __pos__(self) -> ndarray: ...
+    def __abs__(self) -> ndarray: ...
+    def __invert__(self) -> ndarray: ...
+
+    # TODO(shoyer): remove when all methods are defined
     def __getattr__(self, name) -> Any: ...
 
 
 def array(
     object: object,
-    dtype: dtype = ...,
+    dtype: _DtypeLike = ...,
     copy: bool = ...,
     subok: bool = ...,
     ndmin: int = ...) -> ndarray: ...
+
+
+# TODO(shoyer): remove when the full numpy namespace is defined
+def __getattr__(name: str) -> Any: ...

tests/test_simple.py

@@ -1,5 +1,134 @@
+"""Simple expression that should pass with mypy."""
+import operator
+
 import numpy as np
+from typing import Iterable
+
+# Basic checks
+array = np.array([1, 2])
+def ndarray_func(x: np.ndarray) -> np.ndarray:
+    return x
+ndarray_func(np.array([1, 2]))
+array == 1
+array.dtype == float
+
+# Dtype construction
+np.dtype(float)
+np.dtype(np.float64)
+np.dtype('float64')
+np.dtype(np.dtype(float))
+np.dtype(('U', 10))
+np.dtype((np.int32, (2, 2)))
+np.dtype([('R', 'u1'), ('G', 'u1'), ('B', 'u1')])
+np.dtype([('R', 'u1', (2, 2))])
+np.dtype({'col1': ('U10', 0), 'col2': ('float32', 10)})
+np.dtype((np.int32, {'real': (np.int16, 0), 'imag': (np.int16, 2)}))
+np.dtype((np.int32, (np.int8, 4)))
+
+# Iteration and indexing
+def iterable_func(x: Iterable) -> Iterable:
+    return x
+iterable_func(array)
+[element for element in array]
+iter(array)
+zip(array, array)
+array[1]
+array[:]
+array[...]
+array[:] = 0
+
+array_2d = np.ones((3, 3))
+array_2d[:2, :2]
+array_2d[..., 0]
+array_2d[:2, :2] = 0
+
+# Other special methods
+len(array)
+str(array)
+array_scalar = np.array(1)
+int(array_scalar)
+float(array_scalar)
+# currently does not work due to https://github.com/python/typeshed/issues/1904
+# complex(array_scalar)
+bytes(array_scalar)
+operator.index(array_scalar)
+bool(array_scalar)
+
+# comparisons
+array < 1
+array <= 1
+array == 1
+array != 1
+array > 1
+array >= 1
+1 < array
+1 <= array
+1 == array
+1 != array
+1 > array
+1 >= array
+
+# binary arithmetic
+array + 1
+1 + array
+array += 1
+
+array - 1
+1 - array
+array -= 1
+
+array * 1
+1 * array
+array *= 1
+
+array / 1
+1 / array
+array /= 1
+
+array // 1
+1 // array
+array //= 1
+
+array % 1
+1 % array
+array %= 1
+
+# Overloading divmod() is not yet support in typeshed:
+# https://github.com/python/typing/issues/541
+# divmod(array, 1)
+# divmod(1, array)
+
+array ** 1
+1 ** array
+array **= 1
+
+array << 1
+1 << array
+array <<= 1
+
+array >> 1
+1 >> array
+array >>= 1
+
+array & 1
+1 & array
+array &= 1
+
+array ^ 1
+1 ^ array
+array ^= 1
+
+array | 1
+1 | array
+array |= 1
+
+array @ array
 
-def foo(a: np.ndarray): pass
+# unary arithmetic
+-array
++array
+abs(array)
+~array
 
-foo(np.array(1))
+# Other methods
+np.array([1, 2]).transpose()