f

Author: crusaderky

src/array_api_extra/_lib/_at.py

@@ -6,7 +6,6 @@
 import operator
 from collections.abc import Callable
 from enum import Enum
-from numbers import Number
 from types import ModuleType
 from typing import ClassVar, cast
 
@@ -189,7 +188,7 @@ def __getitem__(self, idx: Index, /) -> at:  # numpydoc ignore=PR01,RT01
     def _update_common(
         self,
         at_op: _AtOp,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None,
         xp: ModuleType | None,
@@ -254,7 +253,7 @@ def _update_common(
 
     def set(
         self,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -270,8 +269,8 @@ def set(
     def _iop(
         self,
         at_op: _AtOp,
-        elwise_op: Callable[[Array, Array | Number], Array],
-        y: Array | Number,
+        elwise_op: Callable[[Array, Array | object], Array],
+        y: Array | object,
         /,
         copy: bool | None,
         xp: ModuleType | None,
@@ -295,7 +294,7 @@ def _iop(
 
     def add(
         self,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -309,7 +308,7 @@ def add(
 
     def subtract(
         self,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -319,7 +318,7 @@ def subtract(
 
     def multiply(
         self,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -329,7 +328,7 @@ def multiply(
 
     def divide(
         self,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -339,7 +338,7 @@ def divide(
 
     def power(
         self,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -349,7 +348,7 @@ def power(
 
     def min(
         self,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,
@@ -362,7 +361,7 @@ def min(
 
     def max(
         self,
-        y: Array | Number,
+        y: Array | object,
         /,
         copy: bool | None = None,
         xp: ModuleType | None = None,

tests/test_at.py

@@ -1,7 +1,6 @@
 import pickle
 from collections.abc import Callable, Generator
 from contextlib import contextmanager
-from numbers import Number
 from types import ModuleType
 from typing import cast
 
@@ -21,7 +20,7 @@ def at_op(
     x: Array,
     idx: Index,
     op: _AtOp,
-    y: Array | Number,
+    y: Array | object,
     copy: bool | None = None,
     xp: ModuleType | None = None,
 ) -> Array:
@@ -43,7 +42,7 @@ def _at_op(
     idx: Index | None,
     idx_pickle: bytes | None,
     op: _AtOp,
-    y: Array | Number,
+    y: Array | object,
     copy: bool | None = None,
     xp: ModuleType | None = None,
 ) -> Array: