Ensure stubs are valid for Python 2 and fix running of tests

The stubs contained an unconditional reference to SupportsBytes, which only exists in Python 3. To make these valid on Python 2, conditionally import that Protocol in Python 3 and otherwise use a dummy class in Python 2. Also have `ndarray` extend `Contains`, while we're here.

This also extends the test suites to run all tests against both Python 2 and Python 3, with the ability to specify that certain tests should only be run against Python 3 (eg to test Python 3 exclusive operators). This should help prevent errors like this moving forward.

One downside of this is that flake8 doesn't understand the `# type:` comments, so it thinks that imports from `typing` are unused. A workaround for this is to add `# noqa: F401` at the end of the relevant imports, though this is a bit tedious.

Finally, change how test requirements are installed and how the `numpy-stubs` package is exposed to mypy, and update the README/Travis file to reflect this. See python/mypy#5007 for more details about the rational behind this change.

.travis.yml

@@ -6,7 +6,6 @@ notifications:
 
 install:
 - pip install -r test-requirements.txt
-- pip install .
 
 script:
 - flake8

numpy-stubs/__init__.pyi

@@ -1,14 +1,32 @@
 import builtins
+import sys
 
+from numpy.core._internal import _ctypes
 from typing import (
-    Any, Dict, Iterable, List, Optional, Mapping, Sequence, Sized,
-    SupportsInt, SupportsFloat, SupportsComplex, SupportsBytes, SupportsAbs,
-    Text, Tuple, Type, TypeVar, Union,
+    Any,
+    Container,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Sized,
+    SupportsAbs,
+    SupportsComplex,
+    SupportsFloat,
+    SupportsInt,
+    Text,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
 )
 
-import sys
-
-from numpy.core._internal import _ctypes
+if sys.version_info[0] < 3:
+    class SupportsBytes: ...
+else:
+    from typing import SupportsBytes
 
 _Shape = Tuple[int, ...]
 
@@ -325,7 +343,7 @@ class _ArrayOrScalarCommon(SupportsInt, SupportsFloat, SupportsComplex,
     def __getattr__(self, name) -> Any: ...
 
 
-class ndarray(_ArrayOrScalarCommon, Iterable, Sized):
+class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container):
     real: ndarray
     imag: ndarray
 

test-requirements.txt

@@ -2,3 +2,6 @@ flake8==3.3.0
 flake8-pyi==17.3.0
 pytest==3.4.2
 mypy==0.590
+# This makes sure that the repo's stubs are accessible. Using MYPYPATH won't
+# work. See https://github.com/python/mypy/issues/5007 for more details.
+.

tests/README.md

@@ -22,6 +22,9 @@ reveal_type(x)   # E: <type name>
 Right now, the error messages and types are must be **contained within
 corresponding mypy message**.
 
+Test files that end in `_py3.py` will only be type checked against Python 3.
+All other test files must be valid in both Python 2 and Python 3.
+
 ## Running the tests
 
 We use `py.test` to orchestrate our tests. You can just run:
@@ -34,6 +37,23 @@ 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>
+mypy <file_path>
+```
+
+Note that for either of these commands, you must run:
+
+```
+pip install -r test-requirements.txt
+```
+
+for the version of python that you're going to be running `py.test` or `mypy`
+with. To ensure you're using the intended version of Python you can use
+`python -m` versions of these commands instead:
+
+```
+python -m pytest
+python -m mypy <file_path>
+python -m pip install -r test-requirements.txt
 ```
+Due to how mypy reads type information in PEP 561 packages, you'll need
+to re-run the `pip install` command each time you change the stubs.

tests/pass/simple.py

@@ -2,11 +2,12 @@
 import operator
 
 import numpy as np
-from typing import Iterable
+from typing import Iterable  # noqa: F401
 
 # Basic checks
 array = np.array([1, 2])
-def ndarray_func(x: np.ndarray) -> np.ndarray:
+def ndarray_func(x):
+    # type: (np.ndarray) -> np.ndarray
     return x
 ndarray_func(np.array([1, 2]))
 array == 1
@@ -28,7 +29,8 @@ def ndarray_func(x: np.ndarray) -> np.ndarray:
 np.dtype((np.int32, (np.int8, 4)))
 
 # Iteration and indexing
-def iterable_func(x: Iterable) -> Iterable:
+def iterable_func(x):
+    # type: (Iterable) -> Iterable
     return x
 iterable_func(array)
 [element for element in array]
@@ -122,8 +124,6 @@ def iterable_func(x: Iterable) -> Iterable:
 1 | array
 array |= 1
 
-array @ array
-
 # unary arithmetic
 -array
 +array

tests/pass/simple_py3.py

@@ -0,0 +1,6 @@
+import numpy as np
+
+array = np.array([1, 2])
+
+# The @ operator is not in python 2
+array @ array

tests/test_stubs.py

@@ -3,37 +3,48 @@
 import pytest
 from mypy import api
 
-ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
-PASS_DIR = os.path.join(os.path.dirname(__file__), "pass")
-FAIL_DIR = os.path.join(os.path.dirname(__file__), "fail")
-REVEAL_DIR = os.path.join(os.path.dirname(__file__), "reveal")
-
-os.environ['MYPYPATH'] = ROOT_DIR
+TESTS_DIR = os.path.dirname(__file__)
+PASS_DIR = os.path.join(TESTS_DIR, "pass")
+FAIL_DIR = os.path.join(TESTS_DIR, "fail")
+REVEAL_DIR = os.path.join(TESTS_DIR, "reveal")
 
 
 def get_test_cases(directory):
     for root, __, files in os.walk(directory):
         for fname in files:
             if os.path.splitext(fname)[-1] == ".py":
-                # yield relative path for nice py.test name
-                yield os.path.relpath(
-                    os.path.join(root, fname), start=directory)
-
-
-@pytest.mark.parametrize("path", get_test_cases(PASS_DIR))
-def test_success(path):
-    stdout, stderr, exitcode = api.run([os.path.join(PASS_DIR, path)])
+                fullpath = os.path.join(root, fname)
+                # Use relative path for nice py.test name
+                relpath = os.path.relpath(fullpath, start=directory)
+                skip_py2 = fname.endswith("_py3.py")
+
+                for py_version_number in (2, 3):
+                    if py_version_number == 2 and skip_py2:
+                        continue
+                    py2_arg = ['--py2'] if py_version_number == 2 else []
+
+                    yield pytest.param(
+                        fullpath,
+                        py2_arg,
+                        # Manually specify a name for the test
+                        id="{} - python{}".format(relpath, py_version_number),
+                    )
+
+
+@pytest.mark.parametrize("path,py2_arg", get_test_cases(PASS_DIR))
+def test_success(path, py2_arg):
+    stdout, stderr, exitcode = api.run([path] + py2_arg)
     assert stdout == ''
     assert exitcode == 0
 
 
-@pytest.mark.parametrize("path", get_test_cases(FAIL_DIR))
-def test_fail(path):
-    stdout, stderr, exitcode = api.run([os.path.join(FAIL_DIR, path)])
+@pytest.mark.parametrize("path,py2_arg", get_test_cases(FAIL_DIR))
+def test_fail(path, py2_arg):
+    stdout, stderr, exitcode = api.run([path] + py2_arg)
 
     assert exitcode != 0
 
-    with open(os.path.join(FAIL_DIR, path)) as fin:
+    with open(path) as fin:
         lines = fin.readlines()
 
     errors = {}
@@ -59,11 +70,11 @@ def test_fail(path):
             pytest.fail(f'Error {repr(errors[lineno])} not found')
 
 
-@pytest.mark.parametrize("path", get_test_cases(REVEAL_DIR))
-def test_reveal(path):
-    stdout, stderr, exitcode = api.run([os.path.join(REVEAL_DIR, path)])
+@pytest.mark.parametrize("path,py2_arg", get_test_cases(REVEAL_DIR))
+def test_reveal(path, py2_arg):
+    stdout, stderr, exitcode = api.run([path] + py2_arg)
 
-    with open(os.path.join(REVEAL_DIR, path)) as fin:
+    with open(path) as fin:
         lines = fin.readlines()
 
     for error_line in stdout.split("\n"):