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__.
         """
@@ -728,7 +728,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__.
         """
@@ -783,7 +783,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__.
         """
@@ -807,7 +807,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__.
         """
@@ -832,7 +832,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__.
         """
@@ -844,7 +844,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__.
         """
@@ -856,7 +856,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__.
         """
@@ -877,7 +877,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__.
         """
@@ -898,7 +898,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__.
         """
@@ -942,7 +942,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__.
         """
@@ -956,7 +956,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__.
         """
@@ -968,7 +968,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__.
         """
@@ -980,7 +980,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__.
         """
@@ -991,7 +991,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__.
         """
@@ -1003,7 +1003,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__.
         """
@@ -1014,7 +1014,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__.
         """
@@ -1026,7 +1026,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__.
         """
@@ -1037,7 +1037,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__.
         """
@@ -1098,7 +1098,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__.
         """
@@ -1108,7 +1108,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__.
         """
@@ -1120,7 +1120,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__.
         """
@@ -1130,7 +1130,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__.
         """
@@ -1142,7 +1142,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__.
         """
@@ -1152,7 +1152,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__.
         """
@@ -1164,7 +1164,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__.
         """
@@ -1174,7 +1174,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__.
         """
@@ -1209,7 +1209,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__.
         """
@@ -1219,7 +1219,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__.
         """
@@ -1231,7 +1231,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__.
         """
@@ -1241,7 +1241,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__.
         """
@@ -1253,7 +1253,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__.
         """
@@ -1263,7 +1263,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__.
         """