Define special methods for ndarray and add more extensive tests.

numpy/__init__.pyi

@@ -1,15 +1,66 @@
-# 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, Sequence, Sized,
+    SupportsInt, SupportsFloat, SupportsComplex, SupportsBytes, SupportsAbs,
+    Text, Tuple, Union,
+)
+
+import sys
 
 from numpy.core._internal import _ctypes
 
 _Shape = Tuple[int, ...]
 
+# Anything that can be coerced to a shape tuple
+_ShapeLike = Union[int, Sequence[int]]
+
+_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[
+    dtype,
+    # default data type (float64)
+    None,
+    # array-scalar types and generic types
+    type,  # TODO: enumerate these when we add type hints for numpy scalars
+    # TODO: add a protocol for anything with a dtype attribute
+    # character codes, type strings or comma-separated fields, e.g., 'float64'
+    str,
+    # (flexible_dtype, itemsize)
+    Tuple[_DtypeLikeNested, int],
+    # (fixed_dtype, shape)
+    Tuple[_DtypeLikeNested, _ShapeLike],
+    # [(field_name, field_dtype, field_shape), ...]
+    List[Union[
+        Tuple[Union[str, Tuple[str, str]], _DtypeLikeNested],
+        Tuple[Union[str, Tuple[str, str]], _DtypeLikeNested, _ShapeLike]]],
+    # {'names': ..., 'formats': ..., 'offsets': ..., 'titles': ...,
+    #  'itemsize': ...}
+    # TODO: use TypedDict when/if it's officially supported
+    Dict[str, Union[
+        Sequence[str],  # names
+        Sequence[_DtypeLikeNested],  # formats
+        Sequence[int],  # offsets
+        Sequence[Union[bytes, Text, None]],  # titles
+        int,  # itemsize
+    ]],
+    # {'field1': ..., 'field2': ..., ...}
+    Dict[str, Tuple[_DtypeLikeNested, int]],
+    # (base_dtype, new_dtype)
+    Tuple[_DtypeLikeNested, _DtypeLikeNested],
+]
+
+
 class dtype:
     names: Optional[Tuple[str, ...]]
 
+    def __init__(self,
+                 obj: _DtypeLike,
+                 align: bool = ...,
+                 copy: bool = ...) -> None: ...
+
     @property
     def alignment(self) -> int: ...
 
@@ -83,7 +134,7 @@ class dtype:
     def type(self) -> builtins.type: ...
 
 
-_dtype_class = dtype      # for ndarray type
+_Dtype = dtype  # to avoid name conflicts with ndarray.dtype
 
 
 class _flagsobj:
@@ -144,19 +195,23 @@ class flatiter:
     def __next__(self) -> Any: ...
 
 
-class ndarray:
-    dtype: _dtype_class
+class ndarray(Iterable, Sized, SupportsInt, SupportsFloat, SupportsComplex,
+              SupportsBytes, SupportsAbs[Any]):
+
     imag: ndarray
     real: ndarray
-    shape: _Shape
-    strides: Tuple[int, ...]
 
     @property
     def T(self) -> ndarray: ...
 
     @property
     def base(self) -> Optional[ndarray]: ...
 
+    @property
+    def dtype(self) -> _Dtype: ...
+    @dtype.setter
+    def dtype(self, value: _DtypeLike): ...
+
     @property
     def ctypes(self) -> _ctypes: ...
 
@@ -181,12 +236,135 @@ class ndarray:
     @property
     def ndim(self) -> int: ...
 
+    @property
+    def shape(self) -> _Shape: ...
+    @shape.setter
+    def shape(self, value: _ShapeLike): ...
+
+    @property
+    def strides(self) -> _Shape: ...
+    @strides.setter
+    def strides(self, value: _ShapeLike): ...
+
+    # 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[0] < 3:
+        def __oct__(self) -> str: ...
+        def __hex__(self) -> str: ...
+        def __nonzero__(self) -> bool: ...
+        def __unicode__(self) -> Text: ...
+    else:
+        def __bool__(self) -> bool: ...
+        def __bytes__(self) -> bytes: ...
+    def __str__(self) -> str: ...
+    def __repr__(self) -> str: ...
+
+    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[0] < 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[:2] >= (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: ...

test-requirements.txt

@@ -1,3 +1,3 @@
 flake8==3.3.0
 
-mypy==0.560.0
+mypy==0.570.0

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(None)
+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', 1)])
+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
+
+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()