Add NumPy scalar hierarchy (#14)

* Add NumPy scalar hierarchy

* Include timedelta64

* Factor out _ArrayLike base class

* Cluster away platform specific types

* fixup! Factor out _ArrayLike base class

* Add _is_real type to fix real and imag typing

* fixup! fixup! Factor out _ArrayLike base class

* s/_is_real/_real_generic/g

* s/_ArrayLike/_ArrayOrScalarCommon

array_like means "can be converted to array"

* Update dtype.type

* Add testing framework

* Fix type annotations

* Update Travis

* Use pytest.fail

Author: Alan Du

.gitignore

@@ -1 +1,3 @@
 .mypy_cache
+.pytest_cache
+__pycache__

.travis.yml

@@ -6,11 +6,10 @@ notifications:
 
 install:
 - pip install -r test-requirements.txt
-- if [[ $TRAVIS_PYTHON_VERSION == '3.6' ]]; then pip install flake8-pyi==17.3.0; fi
 
 script:
 - flake8
-- MYPYPATH="." mypy tests/test_simple.py
+- py.test
 
 cache:
   directories:

numpy/__init__.pyi

@@ -3,7 +3,7 @@ import builtins
 from typing import (
     Any, Dict, Iterable, List, Optional, Mapping, Sequence, Sized,
     SupportsInt, SupportsFloat, SupportsComplex, SupportsBytes, SupportsAbs,
-    Text, Tuple, Union,
+    Text, Tuple, Type, TypeVar, Union,
 )
 
 import sys
@@ -131,7 +131,7 @@ class dtype:
     def str(self) -> builtins.str: ...
 
     @property
-    def type(self) -> builtins.type: ...
+    def type(self) -> Type[generic]: ...
 
 
 _Dtype = dtype  # to avoid name conflicts with ndarray.dtype
@@ -195,35 +195,24 @@ class flatiter:
     def __next__(self) -> Any: ...
 
 
-class ndarray(Iterable, Sized, SupportsInt, SupportsFloat, SupportsComplex,
-              SupportsBytes, SupportsAbs[Any]):
-
-    imag: ndarray
-    real: ndarray
-
+_ArraySelf = TypeVar("_ArraySelf", bound=_ArrayOrScalarCommon)
+class _ArrayOrScalarCommon(SupportsInt, SupportsFloat, SupportsComplex,
+                           SupportsBytes, SupportsAbs[Any]):
     @property
-    def T(self) -> ndarray: ...
+    def T(self: _ArraySelf) -> _ArraySelf: ...
 
     @property
     def base(self) -> Optional[ndarray]: ...
 
     @property
     def dtype(self) -> _Dtype: ...
-    @dtype.setter
-    def dtype(self, value: _DtypeLike): ...
-
-    @property
-    def ctypes(self) -> _ctypes: ...
 
     @property
     def data(self) -> memoryview: ...
 
     @property
     def flags(self) -> _flagsobj: ...
 
-    @property
-    def flat(self) -> flatiter: ...
-
     @property
     def size(self) -> int: ...
 
@@ -238,22 +227,9 @@ class ndarray(Iterable, Sized, SupportsInt, SupportsFloat, SupportsComplex,
 
     @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: ...
@@ -269,17 +245,8 @@ class ndarray(Iterable, Sized, SupportsInt, SupportsFloat, SupportsComplex,
     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 __copy__(self: _ArraySelf, order: str = ...) -> _ArraySelf: ...
+    def __deepcopy__(self: _ArraySelf, memo: dict) -> _ArraySelf: ...
 
     def __lt__(self, other): ...
     def __le__(self, other): ...
@@ -349,15 +316,122 @@ class ndarray(Iterable, Sized, SupportsInt, SupportsFloat, SupportsComplex,
         def __matmul__(self, other): ...
         def __rmatmul__(self, other): ...
 
-    def __neg__(self) -> ndarray: ...
-    def __pos__(self) -> ndarray: ...
-    def __abs__(self) -> ndarray: ...
-    def __invert__(self) -> ndarray: ...
+    def __neg__(self: _ArraySelf) -> _ArraySelf: ...
+    def __pos__(self: _ArraySelf) -> _ArraySelf: ...
+    def __abs__(self: _ArraySelf) -> _ArraySelf: ...
+    def __invert__(self: _ArraySelf) -> _ArraySelf: ...
 
     # TODO(shoyer): remove when all methods are defined
     def __getattr__(self, name) -> Any: ...
 
 
+class ndarray(_ArrayOrScalarCommon, Iterable, Sized):
+    real: ndarray
+    imag: ndarray
+
+    @property
+    def dtype(self) -> _Dtype: ...
+    @dtype.setter
+    def dtype(self, value: _DtypeLike): ...
+
+    @property
+    def ctypes(self) -> _ctypes: ...
+
+    @property
+    def shape(self) -> _Shape: ...
+    @shape.setter
+    def shape(self, value: _ShapeLike): ...
+
+    @property
+    def flat(self) -> flatiter: ...
+
+    @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 __index__(self) -> int: ...
+
+class generic(_ArrayOrScalarCommon):
+    def __init__(self, value: Any = ...) -> None: ...
+    @property
+    def base(self) -> None: ...
+
+class _real_generic(generic):
+    @property
+    def real(self: _ArraySelf) -> _ArraySelf: ...
+
+    @property
+    def imag(self: _ArraySelf) -> _ArraySelf: ...
+
+class number(generic):
+    def __init__(
+        self, value: Union[SupportsInt, SupportsFloat] = ...
+    ) -> None: ...
+class bool_(_real_generic): ...
+class object_(generic): ...
+class datetime64(_real_generic): ...
+
+class integer(number, _real_generic): ...
+class signedinteger(integer): ...
+class int8(signedinteger): ...
+class int16(signedinteger): ...
+class int32(signedinteger): ...
+class int64(signedinteger): ...
+class timedelta64(signedinteger): ...
+
+class unsignedinteger(integer): ...
+class uint8(unsignedinteger): ...
+class uint16(unsignedinteger): ...
+class uint32(unsignedinteger): ...
+class uint64(unsignedinteger): ...
+
+class inexact(number): ...
+class floating(inexact, _real_generic): ...
+class float16(floating): ...
+class float32(floating): ...
+class float64(floating): ...
+
+class complexfloating(inexact):
+    def __init__(
+        self,
+        value: Union[SupportsInt, SupportsFloat, SupportsComplex,
+                     complex] = ...,
+    ) -> None: ...
+class complex64(complexfloating):
+    @property
+    def real(self) -> float32: ...
+    @property
+    def imag(self) -> float32: ...
+class complex128(complexfloating):
+    @property
+    def real(self) -> float64: ...
+    @property
+    def imag(self) -> float64: ...
+
+class flexible(_real_generic): ...
+class void(flexible): ...
+class character(_real_generic): ...
+class bytes_(character): ...
+class str_(character): ...
+
+# TODO(alan): Platform dependent types
+# longcomplex, longdouble, longfloat
+# bytes, short, intc, intp, longlong
+# half, single, double, longdouble
+# uint_, int_, float_, complex_
+# float128, complex256
+# float96
+
+
 def array(
     object: object,
     dtype: _DtypeLike = ...,

setup.cfg

@@ -7,3 +7,4 @@
 
 [flake8]
 ignore = E301, E302, E305, E701, E704
+exclude = .git,__pycache__,tests/reveal

test-requirements.txt

@@ -1,3 +1,4 @@
 flake8==3.3.0
-
-mypy==0.570.0
+flake8-pyi==17.3.0
+pytest==3.4.2
+mypy==0.580.0

tests/README.md

@@ -1,9 +1,39 @@
 Testing
 =======
 
-To run these tests:
+There are three main directories of tests right now:
 
-    export MYPYPATH='..'
-    mypy test_simple.py
+- `pass/` which contain Python files that must pass `mypy` checking with
+  no type errors
+- `fail/` which contain Python files that must *fail* `mypy` checking
+  with the annotated errors
+- `reveal/` which contain Python files that must output the correct
+  types with `reveal_type`
 
-In future, this should change to use the test framework used by mypy.
+`fail` and `reveal` are annotated with comments that specify what error
+`mypy` threw and what type should be revealed respectively. The format
+looks like:
+
+```python
+bad_function   # E: <error message>
+reveal_type(x)   # E: <type name>
+```
+
+Right now, the error messages and types are must be **contained within
+corresponding mypy message**.
+
+## Running the tests
+
+We use `py.test` to orchestrate our tests. You can just run:
+
+```
+py.test
+```
+
+to run the entire test suite. To run `mypy` on a specific file (which
+can be useful for debugging), you can also run:
+
+```
+$ cd tests
+$ MYPYPATH=.. mypy <file_path>
+```

tests/fail/scalars.py

@@ -0,0 +1,9 @@
+import numpy as np
+
+# Construction
+
+np.float32(3j)  # E: incompatible type
+np.complex64([])  # E: incompatible type
+
+np.complex64(1, 2)  # E: Too many arguments
+# TODO: protocols (can't check for non-existent protocols w/ __getattr__)

tests/pass/scalars.py

@@ -0,0 +1,52 @@
+import numpy as np
+
+
+# Construction
+class C:
+    def __complex__(self):
+        return 3j
+
+
+class A:
+    def __float__(self):
+        return 4
+
+
+np.complex32(3)
+np.complex64(3j)
+np.complex64(C())
+
+np.int8(4)
+np.int16(3.4)
+np.int32(4)
+np.int64(-1)
+np.uint8(A())
+np.uint32()
+
+np.float16(A())
+np.float32(16)
+np.float64(3.0)
+
+np.bytes_(b"hello")
+np.str_("hello")
+
+# Protocols
+float(np.int8(4))
+int(np.int16(5))
+np.int8(np.float32(6))
+
+# TODO(alan): test after https://github.com/python/typeshed/pull/2004
+# complex(np.int32(8))
+
+abs(np.int8(4))
+
+# Array-ish semantics
+np.int8().real
+np.int16().imag
+np.int32().data
+np.int64().flags
+
+np.uint8().itemsize * 2
+np.uint16().ndim + 1
+np.uint32().strides
+np.uint64().shape

tests/test_simple.py renamed to tests/pass/simple.py

@@ -0,0 +1,13 @@
+import numpy as np
+
+x = np.complex64(3 + 2j)
+
+reveal_type(x.real)  # E: numpy.float32
+reveal_type(x.imag)  # E: numpy.float32
+
+reveal_type(x.real.real)  # E: numpy.float32
+reveal_type(x.real.imag)  # E: numpy.float32
+
+reveal_type(x.itemsize)   # E: int
+reveal_type(x.shape)      # E: tuple[builtins.int]
+reveal_type(x.strides)    # E: tuple[builtins.int]