Skip to content
Open
22 changes: 18 additions & 4 deletions src/array_api_stubs/_draft/array_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

__all__ = ["array"]

from typing import SupportsIndex
from ._types import (
array,
dtype as Dtype,
Expand Down Expand Up @@ -608,7 +609,8 @@ def __getitem__(
slice,
ellipsis,
None,
Tuple[Union[int, slice, ellipsis, None], ...],
SupportsIndex,
Tuple[Union[int, slice, ellipsis, None, SupportsIndex], ...],
array,
],
/,
Expand All @@ -620,9 +622,13 @@ def __getitem__(
----------
self: array
array instance.
key: Union[int, slice, ellipsis, None, Tuple[Union[int, slice, ellipsis, None], ...], array]
key: Union[int, slice, ellipsis, None, SupportsIndex, Tuple[Union[int, slice, ellipsis, None, SupportsIndex], ...], array]
index key.


.. note::
``key`` can only be an array if it is valid for boolean array indexing, or when it supports ``__index__()`` as a 0-dimensional integer array.

Returns
-------
out: array
Expand Down Expand Up @@ -1077,7 +1083,12 @@ def __rshift__(self: array, other: Union[int, array], /) -> array:
def __setitem__(
self: array,
key: Union[
int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], array
int,
slice,
ellipsis,
SupportsIndex,
Tuple[Union[int, slice, ellipsis, SupportsIndex], ...],
array,
],
value: Union[int, float, bool, array],
/,
Expand All @@ -1089,12 +1100,15 @@ def __setitem__(
----------
self: array
array instance.
key: Union[int, slice, ellipsis, Tuple[Union[int, 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`).


.. note::
``key`` can only be an array if it is valid for boolean array indexing, or when it supports ``__index__()`` as a 0-dimensional integer array.

.. note::

Setting array values must not affect the data type of ``self``.
Expand Down