Skip to content

Commit 753697f

Browse files
committed
Replace int with SupportsIndex in indexing methods hints
1 parent 630149c commit 753697f

File tree

1 file changed

+12
-5
lines changed

1 file changed

+12
-5
lines changed

Diff for: src/array_api_stubs/_draft/array_object.py

+12-5
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
__all__ = ["array"]
44

5+
from typing import SupportsIndex
56
from ._types import (
67
array,
78
dtype as Dtype,
@@ -604,11 +605,11 @@ def __ge__(self: array, other: Union[int, float, array], /) -> array:
604605
def __getitem__(
605606
self: array,
606607
key: Union[
607-
int,
608+
SupportsIndex,
608609
slice,
609610
ellipsis,
610611
None,
611-
Tuple[Union[int, slice, ellipsis, None], ...],
612+
Tuple[Union[SupportsIndex, slice, ellipsis, None], ...],
612613
array,
613614
],
614615
/,
@@ -620,9 +621,13 @@ def __getitem__(
620621
----------
621622
self: array
622623
array instance.
623-
key: Union[int, slice, ellipsis, None, Tuple[Union[int, slice, ellipsis, None], ...], array]
624+
key: Union[SupportsIndex, slice, ellipsis, None, Tuple[Union[SupportsIndex, slice, ellipsis, None], ...], array]
624625
index key.
625626
627+
628+
.. note::
629+
``key`` can only be an array if it is valid for boolean array indexing, or supports ``__index__()``.
630+
626631
Returns
627632
-------
628633
out: array
@@ -1077,7 +1082,7 @@ def __rshift__(self: array, other: Union[int, array], /) -> array:
10771082
def __setitem__(
10781083
self: array,
10791084
key: Union[
1080-
int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], array
1085+
SupportsIndex, slice, ellipsis, Tuple[Union[SupportsIndex, slice, ellipsis], ...], array
10811086
],
10821087
value: Union[int, float, bool, array],
10831088
/,
@@ -1089,11 +1094,13 @@ def __setitem__(
10891094
----------
10901095
self: array
10911096
array instance.
1092-
key: Union[int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], array]
1097+
key: Union[SupportsIndex, slice, ellipsis, Tuple[Union[SupportsIndex, slice, ellipsis], ...], array]
10931098
index key.
10941099
value: Union[int, float, bool, array]
10951100
value(s) to set. Must be compatible with ``self[key]`` (see :ref:`broadcasting`).
10961101
1102+
.. note::
1103+
``key`` can only be an array if it is valid for boolean array indexing, or supports ``__index__()``.
10971104
10981105
.. note::
10991106

0 commit comments

Comments
 (0)