verbose Python scalar types

Author: crusaderky

array_api_strict/_array_object.py

@@ -199,7 +199,7 @@ def __array__(
     # NumPy behavior
 
     def _check_allowed_dtypes(
-        self, other: Array | complex, dtype_category: str, op: str
+        self, other: Array | bool | int | float | complex, dtype_category: str, op: str
     ) -> Array:
         """
         Helper function for operators to only allow specific input dtypes
@@ -241,7 +241,7 @@ def _check_allowed_dtypes(
 
         return other
 
-    def _check_device(self, other: Array | complex) -> None:
+    def _check_device(self, other: Array | bool | int | float | complex) -> None:
         """Check that other is on a device compatible with the current array"""
         if isinstance(other, (bool, int, float, complex)):
             return
@@ -252,7 +252,7 @@ def _check_device(self, other: Array | complex) -> None:
             raise TypeError(f"Expected Array | python scalar; got {type(other)}")
 
     # Helper function to match the type promotion rules in the spec
-    def _promote_scalar(self, scalar: complex) -> Array:
+    def _promote_scalar(self, scalar: bool | int | float | complex) -> Array:
         """
         Returns a promoted version of a Python scalar appropriate for use with
         operations on self.
@@ -546,7 +546,7 @@ def __abs__(self) -> Array:
         res = self._array.__abs__()
         return self.__class__._new(res, device=self.device)
 
-    def __add__(self, other: Array | complex, /) -> Array:
+    def __add__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __add__.
         """
@@ -558,7 +558,7 @@ def __add__(self, other: Array | complex, /) -> Array:
         res = self._array.__add__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __and__(self, other: Array | int, /) -> Array:
+    def __and__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __and__.
         """
@@ -655,7 +655,7 @@ def __dlpack_device__(self) -> tuple[IntEnum, int]:
         # Note: device support is required for this
         return self._array.__dlpack_device__()
 
-    def __eq__(self, other: Array | complex, /) -> Array:  # type: ignore[override]
+    def __eq__(self, other: Array | bool | int | float | complex, /) -> Array:  # type: ignore[override]
         """
         Performs the operation __eq__.
         """
@@ -681,7 +681,7 @@ def __float__(self) -> float:
         res = self._array.__float__()
         return res
 
-    def __floordiv__(self, other: Array | float, /) -> Array:
+    def __floordiv__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __floordiv__.
         """
@@ -693,7 +693,7 @@ def __floordiv__(self, other: Array | float, /) -> Array:
         res = self._array.__floordiv__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ge__(self, other: Array | float, /) -> Array:
+    def __ge__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __ge__.
         """
@@ -729,7 +729,7 @@ def __getitem__(
         res = self._array.__getitem__(np_key)
         return self._new(res, device=self.device)
 
-    def __gt__(self, other: Array | float, /) -> Array:
+    def __gt__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __gt__.
         """
@@ -784,7 +784,7 @@ def __iter__(self) -> Iterator[Array]:
         # implemented, which implies iteration on 1-D arrays.
         return (Array._new(i, device=self.device) for i in self._array)
 
-    def __le__(self, other: Array | float, /) -> Array:
+    def __le__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __le__.
         """
@@ -808,7 +808,7 @@ def __lshift__(self, other: Array | int, /) -> Array:
         res = self._array.__lshift__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __lt__(self, other: Array | float, /) -> Array:
+    def __lt__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __lt__.
         """
@@ -833,7 +833,7 @@ def __matmul__(self, other: Array, /) -> Array:
         res = self._array.__matmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __mod__(self, other: Array | float, /) -> Array:
+    def __mod__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __mod__.
         """
@@ -845,7 +845,7 @@ def __mod__(self, other: Array | float, /) -> Array:
         res = self._array.__mod__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __mul__(self, other: Array | complex, /) -> Array:
+    def __mul__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __mul__.
         """
@@ -857,7 +857,7 @@ def __mul__(self, other: Array | complex, /) -> Array:
         res = self._array.__mul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ne__(self, other: Array | complex, /) -> Array:  # type: ignore[override]
+    def __ne__(self, other: Array | bool | int | float | complex, /) -> Array:  # type: ignore[override]
         """
         Performs the operation __ne__.
         """
@@ -878,7 +878,7 @@ def __neg__(self) -> Array:
         res = self._array.__neg__()
         return self.__class__._new(res, device=self.device)
 
-    def __or__(self, other: Array | int, /) -> Array:
+    def __or__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __or__.
         """
@@ -899,7 +899,7 @@ def __pos__(self) -> Array:
         res = self._array.__pos__()
         return self.__class__._new(res, device=self.device)
 
-    def __pow__(self, other: Array | complex, /) -> Array:
+    def __pow__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __pow__.
         """
@@ -936,7 +936,7 @@ def __setitem__(
             | Array
             | tuple[int | slice | EllipsisType, ...]
         ),
-        value: Array | complex,
+        value: Array | bool | int | float | complex,
         /,
     ) -> None:
         """
@@ -949,7 +949,7 @@ def __setitem__(
         np_key = key._array if isinstance(key, Array) else key
         self._array.__setitem__(np_key, asarray(value)._array)
 
-    def __sub__(self, other: Array | complex, /) -> Array:
+    def __sub__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __sub__.
         """
@@ -963,7 +963,7 @@ def __sub__(self, other: Array | complex, /) -> Array:
 
     # PEP 484 requires int to be a subtype of float, but __truediv__ should
     # not accept int.
-    def __truediv__(self, other: Array | complex, /) -> Array:
+    def __truediv__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __truediv__.
         """
@@ -975,7 +975,7 @@ def __truediv__(self, other: Array | complex, /) -> Array:
         res = self._array.__truediv__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __xor__(self, other: Array | int, /) -> Array:
+    def __xor__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __xor__.
         """
@@ -987,7 +987,7 @@ def __xor__(self, other: Array | int, /) -> Array:
         res = self._array.__xor__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __iadd__(self, other: Array | complex, /) -> Array:
+    def __iadd__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __iadd__.
         """
@@ -998,7 +998,7 @@ def __iadd__(self, other: Array | complex, /) -> Array:
         self._array.__iadd__(other._array)
         return self
 
-    def __radd__(self, other: Array | complex, /) -> Array:
+    def __radd__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __radd__.
         """
@@ -1010,7 +1010,7 @@ def __radd__(self, other: Array | complex, /) -> Array:
         res = self._array.__radd__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __iand__(self, other: Array | int, /) -> Array:
+    def __iand__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __iand__.
         """
@@ -1021,7 +1021,7 @@ def __iand__(self, other: Array | int, /) -> Array:
         self._array.__iand__(other._array)
         return self
 
-    def __rand__(self, other: Array | int, /) -> Array:
+    def __rand__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __rand__.
         """
@@ -1033,7 +1033,7 @@ def __rand__(self, other: Array | int, /) -> Array:
         res = self._array.__rand__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ifloordiv__(self, other: Array | float, /) -> Array:
+    def __ifloordiv__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __ifloordiv__.
         """
@@ -1044,7 +1044,7 @@ def __ifloordiv__(self, other: Array | float, /) -> Array:
         self._array.__ifloordiv__(other._array)
         return self
 
-    def __rfloordiv__(self, other: Array | float, /) -> Array:
+    def __rfloordiv__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __rfloordiv__.
         """
@@ -1105,7 +1105,7 @@ def __rmatmul__(self, other: Array, /) -> Array:
         res = self._array.__rmatmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __imod__(self, other: Array | float, /) -> Array:
+    def __imod__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __imod__.
         """
@@ -1115,7 +1115,7 @@ def __imod__(self, other: Array | float, /) -> Array:
         self._array.__imod__(other._array)
         return self
 
-    def __rmod__(self, other: Array | float, /) -> Array:
+    def __rmod__(self, other: Array | int | float, /) -> Array:
         """
         Performs the operation __rmod__.
         """
@@ -1127,7 +1127,7 @@ def __rmod__(self, other: Array | float, /) -> Array:
         res = self._array.__rmod__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __imul__(self, other: Array | complex, /) -> Array:
+    def __imul__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __imul__.
         """
@@ -1137,7 +1137,7 @@ def __imul__(self, other: Array | complex, /) -> Array:
         self._array.__imul__(other._array)
         return self
 
-    def __rmul__(self, other: Array | complex, /) -> Array:
+    def __rmul__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __rmul__.
         """
@@ -1149,7 +1149,7 @@ def __rmul__(self, other: Array | complex, /) -> Array:
         res = self._array.__rmul__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ior__(self, other: Array | int, /) -> Array:
+    def __ior__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __ior__.
         """
@@ -1159,7 +1159,7 @@ def __ior__(self, other: Array | int, /) -> Array:
         self._array.__ior__(other._array)
         return self
 
-    def __ror__(self, other: Array | int, /) -> Array:
+    def __ror__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __ror__.
         """
@@ -1171,7 +1171,7 @@ def __ror__(self, other: Array | int, /) -> Array:
         res = self._array.__ror__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ipow__(self, other: Array | complex, /) -> Array:
+    def __ipow__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __ipow__.
         """
@@ -1181,7 +1181,7 @@ def __ipow__(self, other: Array | complex, /) -> Array:
         self._array.__ipow__(other._array)
         return self
 
-    def __rpow__(self, other: Array | complex, /) -> Array:
+    def __rpow__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __rpow__.
         """
@@ -1216,7 +1216,7 @@ def __rrshift__(self, other: Array | int, /) -> Array:
         res = self._array.__rrshift__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __isub__(self, other: Array | complex, /) -> Array:
+    def __isub__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __isub__.
         """
@@ -1226,7 +1226,7 @@ def __isub__(self, other: Array | complex, /) -> Array:
         self._array.__isub__(other._array)
         return self
 
-    def __rsub__(self, other: Array | complex, /) -> Array:
+    def __rsub__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __rsub__.
         """
@@ -1238,7 +1238,7 @@ def __rsub__(self, other: Array | complex, /) -> Array:
         res = self._array.__rsub__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __itruediv__(self, other: Array | complex, /) -> Array:
+    def __itruediv__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __itruediv__.
         """
@@ -1248,7 +1248,7 @@ def __itruediv__(self, other: Array | complex, /) -> Array:
         self._array.__itruediv__(other._array)
         return self
 
-    def __rtruediv__(self, other: Array | complex, /) -> Array:
+    def __rtruediv__(self, other: Array | int | float | complex, /) -> Array:
         """
         Performs the operation __rtruediv__.
         """
@@ -1260,7 +1260,7 @@ def __rtruediv__(self, other: Array | complex, /) -> Array:
         res = self._array.__rtruediv__(other._array)
         return self.__class__._new(res, device=self.device)
 
-    def __ixor__(self, other: Array | int, /) -> Array:
+    def __ixor__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __ixor__.
         """
@@ -1270,7 +1270,7 @@ def __ixor__(self, other: Array | int, /) -> Array:
         self._array.__ixor__(other._array)
         return self
 
-    def __rxor__(self, other: Array | int, /) -> Array:
+    def __rxor__(self, other: Array | bool | int, /) -> Array:
         """
         Performs the operation __rxor__.
         """