TYP,MAINT: Allow ndindex to accept integer tuples

Author: Bas van Beek

numpy/__init__.pyi

@@ -3314,6 +3314,9 @@ class ndenumerate(Generic[_ScalarType]):
     def __iter__(self: _T) -> _T: ...
 
 class ndindex:
+    @overload
+    def __init__(self, shape: tuple[SupportsIndex, ...], /) -> None: ...
+    @overload
     def __init__(self, *shape: SupportsIndex) -> None: ...
     def __iter__(self: _T) -> _T: ...
     def __next__(self) -> _Shape: ...

numpy/typing/tests/data/fail/index_tricks.pyi

@@ -3,6 +3,7 @@ import numpy as np
 AR_LIKE_i: list[int]
 AR_LIKE_f: list[float]
 
+np.ndindex([1, 2, 3])  # E: No overload variant
 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: No overload variant
 np.mgrid[1]  # E: Invalid index type

numpy/typing/tests/data/reveal/index_tricks.pyi

@@ -24,6 +24,8 @@ reveal_type(iter(np.ndenumerate(AR_i8)))  # E: Iterator[Tuple[builtins.tuple[bui
 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], str_]]
 
+reveal_type(np.ndindex(1, 2, 3))  # E: numpy.ndindex
+reveal_type(np.ndindex((1, 2, 3)))  # E: numpy.ndindex
 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]