index

Author: crusaderky

src/array_api_extra/_lib/_utils/_typing.pyi

@@ -95,21 +95,11 @@ class DType(Protocol):  # pylint: disable=missing-class-docstring
 class Device(Protocol):  # pylint: disable=missing-class-docstring
     pass
 
-GetIndex: TypeAlias = (  # type: ignore[no-any-explicit]
-    int
-    | slice
-    | EllipsisType
-    | Array
-    | None
-    | tuple[int | slice | EllipsisType | None | Array, ...]
-)
 SetIndex: TypeAlias = (  # type: ignore[no-any-explicit]
-    int
-    | slice
-    | EllipsisType
-    | Array
-    | tuple[int | slice | EllipsisType | Array, ...]
-    | GetIndex
+    int | slice | EllipsisType | Array | tuple[int | slice | EllipsisType | Array, ...]
+)
+GetIndex: TypeAlias = (  # type: ignore[no-any-explicit]
+    SetIndex | None | tuple[int | slice | EllipsisType | None | Array, ...]
 )
 
 __all__ = ["Array", "DType", "Device", "GetIndex", "SetIndex"]

tests/test_at.py

@@ -56,7 +56,7 @@ def _at_op(
     """jitted helper of at_op"""
     if idx_pickle:
         idx = pickle.loads(idx_pickle)
-    meth = cast(Callable[..., Array], getattr(at(x, idx), op.value))  # type: ignore[no-any-explicit]
+    meth = cast(Callable[..., Array], getattr(at(x, cast(SetIndex, idx)), op.value))  # type: ignore[no-any-explicit]
     return meth(y, copy=copy, xp=xp)