diff options
author | Bas van Beek <b.f.van.beek@vu.nl> | 2021-02-11 18:21:31 +0100 |
---|---|---|
committer | Bas van Beek <b.f.van.beek@vu.nl> | 2021-02-25 14:06:37 +0100 |
commit | 8a18515f6b5eb97a184b92858741eb4a5aab613a (patch) | |
tree | 80ef24da244a26ed806338a4aa6f5c82ac33ccbf | |
parent | a58784d1df7189944b15cccb8702ae85969b9b02 (diff) | |
download | numpy-8a18515f6b5eb97a184b92858741eb4a5aab613a.tar.gz |
TST: Add typing tests for `np.lib.index_tricks`
-rw-r--r-- | numpy/typing/tests/data/fail/index_tricks.py | 18 | ||||
-rw-r--r-- | numpy/typing/tests/data/pass/index_tricks.py | 62 | ||||
-rw-r--r-- | numpy/typing/tests/data/reveal/index_tricks.py | 61 |
3 files changed, 141 insertions, 0 deletions
diff --git a/numpy/typing/tests/data/fail/index_tricks.py b/numpy/typing/tests/data/fail/index_tricks.py new file mode 100644 index 000000000..706f135b2 --- /dev/null +++ b/numpy/typing/tests/data/fail/index_tricks.py @@ -0,0 +1,18 @@ +from typing import List +import numpy as np + +AR_LIKE_i: List[int] +AR_LIKE_f: List[float] + +np.unravel_index(AR_LIKE_f, (1, 2, 3)) # E: incompatible type +np.ravel_multi_index(AR_LIKE_i, (1, 2, 3), mode="bob") # E: incompatible type +np.mgrid[1] # E: Invalid index type +np.mgrid[...] # E: Invalid index type +np.ogrid[1] # E: Invalid index type +np.ogrid[...] # E: Invalid index type +np.index_exp[1] # E: No overload variant +np.index_exp[...] # E: No overload variant +np.s_[1] # E: cannot be "int" +np.s_[...] # E: cannot be "ellipsis" +np.fill_diagonal(AR_LIKE_f, 2) # E: incompatible type +np.diag_indices(1.0) # E: incompatible type diff --git a/numpy/typing/tests/data/pass/index_tricks.py b/numpy/typing/tests/data/pass/index_tricks.py new file mode 100644 index 000000000..ce7f415f3 --- /dev/null +++ b/numpy/typing/tests/data/pass/index_tricks.py @@ -0,0 +1,62 @@ +from __future__ import annotations +from typing import Any +import numpy as np + +AR_LIKE_b = [[True, True], [True, True]] +AR_LIKE_i = [[1, 2], [3, 4]] +AR_LIKE_f = [[1.0, 2.0], [3.0, 4.0]] +AR_LIKE_U = [["1", "2"], ["3", "4"]] + +AR_i8: np.ndarray[Any, np.dtype[np.int64]] = np.array(AR_LIKE_i, dtype=np.int64) + +np.ndenumerate(AR_i8) +np.ndenumerate(AR_LIKE_f) +np.ndenumerate(AR_LIKE_U) + +np.ndenumerate(AR_i8).iter +np.ndenumerate(AR_LIKE_f).iter +np.ndenumerate(AR_LIKE_U).iter + +next(np.ndenumerate(AR_i8)) +next(np.ndenumerate(AR_LIKE_f)) +next(np.ndenumerate(AR_LIKE_U)) + +iter(np.ndenumerate(AR_i8)) +iter(np.ndenumerate(AR_LIKE_f)) +iter(np.ndenumerate(AR_LIKE_U)) + +iter(np.ndindex(1, 2, 3)) +next(np.ndindex(1, 2, 3)) + +np.unravel_index([22, 41, 37], (7, 6)) +np.unravel_index([31, 41, 13], (7, 6), order='F') +np.unravel_index(1621, (6, 7, 8, 9)) + +np.ravel_multi_index(AR_LIKE_i, (7, 6)) +np.ravel_multi_index(AR_LIKE_i, (7, 6), order='F') +np.ravel_multi_index(AR_LIKE_i, (4, 6), mode='clip') +np.ravel_multi_index(AR_LIKE_i, (4, 4), mode=('clip', 'wrap')) +np.ravel_multi_index((3, 1, 4, 1), (6, 7, 8, 9)) + +np.mgrid[1:1:2] +np.mgrid[1:1:2, None:10] + +np.ogrid[1:1:2] +np.ogrid[1:1:2, None:10] + +np.index_exp[0:1] +np.index_exp[0:1, None:3] + +np.s_[0:1] +np.s_[0:1, None:3] + +np.ix_(AR_LIKE_b[0]) +np.ix_(AR_LIKE_i[0], AR_LIKE_f[0]) +np.ix_(AR_i8[0]) + +np.fill_diagonal(AR_i8, 5) + +np.diag_indices(4) +np.diag_indices(2, 3) + +np.diag_indices_from(AR_i8) diff --git a/numpy/typing/tests/data/reveal/index_tricks.py b/numpy/typing/tests/data/reveal/index_tricks.py new file mode 100644 index 000000000..dc061d314 --- /dev/null +++ b/numpy/typing/tests/data/reveal/index_tricks.py @@ -0,0 +1,61 @@ +from typing import Any, List +import numpy as np + +AR_LIKE_b: List[bool] +AR_LIKE_i: List[int] +AR_LIKE_f: List[float] +AR_LIKE_U: List[str] + +AR_i8: np.ndarray[Any, np.dtype[np.int64]] + +reveal_type(np.ndenumerate(AR_i8)) # E: numpy.ndenumerate[{int64}] +reveal_type(np.ndenumerate(AR_LIKE_f)) # E: numpy.ndenumerate[{double}] +reveal_type(np.ndenumerate(AR_LIKE_U)) # E: numpy.ndenumerate[numpy.str_] + +reveal_type(np.ndenumerate(AR_i8).iter) # E: numpy.flatiter[numpy.ndarray[Any, numpy.dtype[{int64}]]] +reveal_type(np.ndenumerate(AR_LIKE_f).iter) # E: numpy.flatiter[numpy.ndarray[Any, numpy.dtype[{double}]]] +reveal_type(np.ndenumerate(AR_LIKE_U).iter) # E: numpy.flatiter[numpy.ndarray[Any, numpy.dtype[numpy.str_]]] + +reveal_type(next(np.ndenumerate(AR_i8))) # E: Tuple[builtins.tuple[builtins.int], {int64}] +reveal_type(next(np.ndenumerate(AR_LIKE_f))) # E: Tuple[builtins.tuple[builtins.int], {double}] +reveal_type(next(np.ndenumerate(AR_LIKE_U))) # E: Tuple[builtins.tuple[builtins.int], numpy.str_] + +reveal_type(iter(np.ndenumerate(AR_i8))) # E: Iterator[Tuple[builtins.tuple[builtins.int], {int64}]] +reveal_type(iter(np.ndenumerate(AR_LIKE_f))) # E: Iterator[Tuple[builtins.tuple[builtins.int], {double}]] +reveal_type(iter(np.ndenumerate(AR_LIKE_U))) # E: Iterator[Tuple[builtins.tuple[builtins.int], numpy.str_]] + +reveal_type(iter(np.ndindex(1, 2, 3))) # E: Iterator[builtins.tuple[builtins.int]] +reveal_type(next(np.ndindex(1, 2, 3))) # E: builtins.tuple[builtins.int] + +reveal_type(np.unravel_index([22, 41, 37], (7, 6))) # E: tuple[Union[{intp}, numpy.ndarray[Any, numpy.dtype[{intp}]]]] +reveal_type(np.unravel_index([31, 41, 13], (7, 6), order="F")) # E: tuple[Union[{intp}, numpy.ndarray[Any, numpy.dtype[{intp}]]]] +reveal_type(np.unravel_index(1621, (6, 7, 8, 9))) # E: tuple[Union[{intp}, numpy.ndarray[Any, numpy.dtype[{intp}]]]] + +reveal_type(np.ravel_multi_index(AR_LIKE_i, (7, 6))) # E: Union[{intp}, numpy.ndarray[Any, numpy.dtype[{intp}]]] +reveal_type(np.ravel_multi_index(AR_LIKE_i, (7, 6), order="F")) # E: Union[{intp}, numpy.ndarray[Any, numpy.dtype[{intp}]]] +reveal_type(np.ravel_multi_index(AR_LIKE_i, (4, 6), mode="clip")) # E: Union[{intp}, numpy.ndarray[Any, numpy.dtype[{intp}]]] +reveal_type(np.ravel_multi_index(AR_LIKE_i, (4, 4), mode=("clip", "wrap"))) # E: Union[{intp}, numpy.ndarray[Any, numpy.dtype[{intp}]]] +reveal_type(np.ravel_multi_index((3, 1, 4, 1), (6, 7, 8, 9))) # E: Union[{intp}, numpy.ndarray[Any, numpy.dtype[{intp}]]] + +reveal_type(np.mgrid[1:1:2]) # E: numpy.ndarray[Any, numpy.dtype[Any]] +reveal_type(np.mgrid[1:1:2, None:10]) # E: numpy.ndarray[Any, numpy.dtype[Any]] + +reveal_type(np.ogrid[1:1:2]) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]] +reveal_type(np.ogrid[1:1:2, None:10]) # E: list[numpy.ndarray[Any, numpy.dtype[Any]]] + +reveal_type(np.index_exp[0:1]) # E: Tuple[builtins.slice] +reveal_type(np.index_exp[0:1, None:3]) # E: Tuple[builtins.slice, builtins.slice] + +reveal_type(np.s_[0:1]) # E: builtins.slice +reveal_type(np.s_[0:1, None:3]) # E: Tuple[builtins.slice, builtins.slice] + +reveal_type(np.ix_(AR_LIKE_b)) # E: tuple[numpy.ndarray[Any, numpy.dtype[numpy.bool_]]] +reveal_type(np.ix_(AR_LIKE_i, AR_LIKE_f)) # E: tuple[numpy.ndarray[Any, numpy.dtype[{double}]]] +reveal_type(np.ix_(AR_i8)) # E: tuple[numpy.ndarray[Any, numpy.dtype[{int64}]]] + +reveal_type(np.fill_diagonal(AR_i8, 5)) # E: None + +reveal_type(np.diag_indices(4)) # E: tuple[numpy.ndarray[Any, numpy.dtype[{int_}]]] +reveal_type(np.diag_indices(2, 3)) # E: tuple[numpy.ndarray[Any, numpy.dtype[{int_}]]] + +reveal_type(np.diag_indices_from(AR_i8)) # E: tuple[numpy.ndarray[Any, numpy.dtype[{int_}]]] |