Merge pull request #18397 from BvB93/index-tricks

ENH: Add annotations for `np.lib.index_tricks`

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..cbc43fd54
--- /dev/null
+++ b/numpy/typing/tests/data/fail/index_tricks.py
@@ -0,0 +1,14 @@
+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.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..4c4c11959
--- /dev/null
+++ b/numpy/typing/tests/data/pass/index_tricks.py
@@ -0,0 +1,64 @@
+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.index_exp[0, 0:1, ..., [0, 1, 3]]
+
+np.s_[0:1]
+np.s_[0:1, None:3]
+np.s_[0, 0:1, ..., [0, 1, 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..ec2013025
--- /dev/null
+++ b/numpy/typing/tests/data/reveal/index_tricks.py
@@ -0,0 +1,63 @@
+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.index_exp[0, 0:1, ..., [0, 1, 3]])  # E: Tuple[Literal[0]?, builtins.slice, builtins.ellipsis, builtins.list[builtins.int]]
+
+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.s_[0, 0:1, ..., [0, 1, 3]])  # E: Tuple[Literal[0]?, builtins.slice, builtins.ellipsis, builtins.list[builtins.int]]
+
+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_}]]]