From 3c9c047f88164428378dcf2365888231ed073ccc Mon Sep 17 00:00:00 2001 From: Matthew Barber Date: Thu, 21 Mar 2024 10:57:28 +0000 Subject: [PATCH 1/4] Replace `int` with `SupportsIndex` in indexing methods hints --- src/array_api_stubs/_draft/array_object.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/src/array_api_stubs/_draft/array_object.py b/src/array_api_stubs/_draft/array_object.py index 6dd70c278..42e876c98 100644 --- a/src/array_api_stubs/_draft/array_object.py +++ b/src/array_api_stubs/_draft/array_object.py @@ -2,6 +2,7 @@ __all__ = ["array"] +from typing import SupportsIndex from ._types import ( array, dtype as Dtype, @@ -604,11 +605,11 @@ def __ge__(self: array, other: Union[int, float, array], /) -> array: def __getitem__( self: array, key: Union[ - int, + SupportsIndex, slice, ellipsis, None, - Tuple[Union[int, slice, ellipsis, None], ...], + Tuple[Union[SupportsIndex, slice, ellipsis, None], ...], array, ], /, @@ -620,9 +621,13 @@ def __getitem__( ---------- self: array array instance. - key: Union[int, slice, ellipsis, None, Tuple[Union[int, slice, ellipsis, None], ...], array] + key: Union[SupportsIndex, slice, ellipsis, None, Tuple[Union[SupportsIndex, slice, ellipsis, None], ...], array] index key. + + .. note:: + ``key`` can only be an array if it is valid for boolean array indexing, or supports ``__index__()``. + Returns ------- out: array @@ -1077,7 +1082,11 @@ def __rshift__(self: array, other: Union[int, array], /) -> array: def __setitem__( self: array, key: Union[ - int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], array + SupportsIndex, + slice, + ellipsis, + Tuple[Union[SupportsIndex, slice, ellipsis], ...], + array, ], value: Union[int, float, bool, array], /, @@ -1089,11 +1098,13 @@ def __setitem__( ---------- self: array array instance. - key: Union[int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], array] + key: Union[SupportsIndex, slice, ellipsis, Tuple[Union[SupportsIndex, slice, ellipsis], ...], 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 supports ``__index__()``. .. note:: From 18910c077ff549d9dad9004ab75619b89bde00c3 Mon Sep 17 00:00:00 2001 From: Matthew Barber Date: Fri, 22 Mar 2024 08:26:12 +0000 Subject: [PATCH 2/4] Add `int` alongside `SupportIndex` in indexing type hints --- src/array_api_stubs/_draft/array_object.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/array_api_stubs/_draft/array_object.py b/src/array_api_stubs/_draft/array_object.py index 42e876c98..9aafceebc 100644 --- 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`). From ef9c4c1d61306712d801bc883031203f17d941bd Mon Sep 17 00:00:00 2001 From: Matthew Barber Date: Fri, 22 Mar 2024 08:38:09 +0000 Subject: [PATCH 3/4] Add space before `__setitem__()` note in docstring Might fix rendering/warning issue --- src/array_api_stubs/_draft/array_object.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/array_api_stubs/_draft/array_object.py b/src/array_api_stubs/_draft/array_object.py index 9aafceebc..2cc250811 100644 --- a/src/array_api_stubs/_draft/array_object.py +++ b/src/array_api_stubs/_draft/array_object.py @@ -1105,6 +1105,7 @@ def __setitem__( 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 supports ``__index__()``. From 986ba2798dee150a4fe92b79879e9ad3f7acd1fd Mon Sep 17 00:00:00 2001 From: Matthew Barber Date: Wed, 27 Mar 2024 14:03:27 +0000 Subject: [PATCH 4/4] Clarify when `__index__()` is applicable for arrays-as-keys in indexing --- src/array_api_stubs/_draft/array_object.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/array_api_stubs/_draft/array_object.py b/src/array_api_stubs/_draft/array_object.py index 2cc250811..aa9fcc4e0 100644 --- a/src/array_api_stubs/_draft/array_object.py +++ b/src/array_api_stubs/_draft/array_object.py @@ -627,7 +627,7 @@ def __getitem__( .. note:: - ``key`` can only be an array if it is valid for boolean array indexing, or supports ``__index__()``. + ``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 ------- @@ -1107,7 +1107,7 @@ def __setitem__( .. note:: - ``key`` can only be an array if it is valid for boolean array indexing, or supports ``__index__()``. + ``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::