Add int alongside SupportIndex in indexing type hints

data-apis · Mar 22, 2024 · 18910c0 · 18910c0
1 parent 3c9c047
commit 18910c0
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/src/array_api_stubs/_draft/array_object.py b/src/array_api_stubs/_draft/array_object.py
@@ -605,11 +605,12 @@ def __ge__(self: array, other: Union[int, float, array], /) -> array:
     def __getitem__(
         self: array,
         key: Union[
-            SupportsIndex,
+            int,
             slice,
             ellipsis,
             None,
-            Tuple[Union[SupportsIndex, slice, ellipsis, None], ...],
+            SupportsIndex,
+            Tuple[Union[int, slice, ellipsis, None, SupportsIndex], ...],
             array,
         ],
         /,
@@ -621,7 +622,7 @@ def __getitem__(
         ----------
         self: array
             array instance.
-        key: Union[SupportsIndex, slice, ellipsis, None, Tuple[Union[SupportsIndex, slice, ellipsis, None], ...], array]
+        key: Union[int, slice, ellipsis, None, SupportsIndex, Tuple[Union[int, slice, ellipsis, None, SupportsIndex], ...], array]
             index key.
 
 
@@ -1082,10 +1083,11 @@ def __rshift__(self: array, other: Union[int, array], /) -> array:
     def __setitem__(
         self: array,
         key: Union[
-            SupportsIndex,
+            int,
             slice,
             ellipsis,
-            Tuple[Union[SupportsIndex, slice, ellipsis], ...],
+            SupportsIndex,
+            Tuple[Union[int, slice, ellipsis, SupportsIndex], ...],
             array,
         ],
         value: Union[int, float, bool, array],
@@ -1098,7 +1100,7 @@ def __setitem__(
         ----------
         self: array
             array instance.
-        key: Union[SupportsIndex, slice, ellipsis, Tuple[Union[SupportsIndex, slice, ellipsis], ...], array]
+        key: Union[int, slice, ellipsis, SupportsIndex, Tuple[Union[int, slice, ellipsis, SupportsIndex], ...], array]
             index key.
         value: Union[int, float, bool, array]
             value(s) to set. Must be compatible with ``self[key]`` (see :ref:`broadcasting`).