Update numpy.array_api magic methods for complex numbers

Updates from the v2022.12 version of the spec:

- Add __complex__.
- __float__, __int__, and __bool__ are now more lenient in what dtypes they
  can operate on.
- Support complex scalars and dtypes in all operators (except those that
  should not operate on complex numbers).
- Disallow integer scalars that are out of the bounds of the array dtype.
- Update the tests accordingly.

Original NumPy Commit: 8b63fc295bafea3efd3d115964200dbaac7be8a5

Author: asmeurer

array_api_strict/_array_object.py

@@ -24,6 +24,7 @@
     _integer_dtypes,
     _integer_or_boolean_dtypes,
     _floating_dtypes,
+    _complex_floating_dtypes,
     _numeric_dtypes,
     _result_type,
     _dtype_categories,
@@ -139,7 +140,7 @@ def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_categor
 
         if self.dtype not in _dtype_categories[dtype_category]:
             raise TypeError(f"Only {dtype_category} dtypes are allowed in {op}")
-        if isinstance(other, (int, float, bool)):
+        if isinstance(other, (int, complex, float, bool)):
             other = self._promote_scalar(other)
         elif isinstance(other, Array):
             if other.dtype not in _dtype_categories[dtype_category]:
@@ -189,11 +190,23 @@ def _promote_scalar(self, scalar):
                 raise TypeError(
                     "Python int scalars cannot be promoted with bool arrays"
                 )
+            if self.dtype in _integer_dtypes:
+                info = np.iinfo(self.dtype)
+                if not (info.min <= scalar <= info.max):
+                    raise OverflowError(
+                        "Python int scalars must be within the bounds of the dtype for integer arrays"
+                    )
+            # int + array(floating) is allowed
         elif isinstance(scalar, float):
             if self.dtype not in _floating_dtypes:
                 raise TypeError(
                     "Python float scalars can only be promoted with floating-point arrays."
                 )
+        elif isinstance(scalar, complex):
+            if self.dtype not in _complex_floating_dtypes:
+                raise TypeError(
+                    "Python complex scalars can only be promoted with complex floating-point arrays."
+                )
         else:
             raise TypeError("'scalar' must be a Python scalar")
 
@@ -454,11 +467,19 @@ def __bool__(self: Array, /) -> bool:
         # Note: This is an error here.
         if self._array.ndim != 0:
             raise TypeError("bool is only allowed on arrays with 0 dimensions")
-        if self.dtype not in _boolean_dtypes:
-            raise ValueError("bool is only allowed on boolean arrays")
         res = self._array.__bool__()
         return res
 
+    def __complex__(self: Array, /) -> float:
+        """
+        Performs the operation __complex__.
+        """
+        # Note: This is an error here.
+        if self._array.ndim != 0:
+            raise TypeError("complex is only allowed on arrays with 0 dimensions")
+        res = self._array.__complex__()
+        return res
+
     def __dlpack__(self: Array, /, *, stream: None = None) -> PyCapsule:
         """
         Performs the operation __dlpack__.
@@ -492,16 +513,16 @@ def __float__(self: Array, /) -> float:
         # Note: This is an error here.
         if self._array.ndim != 0:
             raise TypeError("float is only allowed on arrays with 0 dimensions")
-        if self.dtype not in _floating_dtypes:
-            raise ValueError("float is only allowed on floating-point arrays")
+        if self.dtype in _complex_floating_dtypes:
+            raise TypeError("float is not allowed on complex floating-point arrays")
         res = self._array.__float__()
         return res
 
     def __floordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __floordiv__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__floordiv__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__floordiv__")
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -512,7 +533,7 @@ def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __ge__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__ge__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__ge__")
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -542,7 +563,7 @@ def __gt__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __gt__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__gt__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__gt__")
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -556,8 +577,8 @@ def __int__(self: Array, /) -> int:
         # Note: This is an error here.
         if self._array.ndim != 0:
             raise TypeError("int is only allowed on arrays with 0 dimensions")
-        if self.dtype not in _integer_dtypes:
-            raise ValueError("int is only allowed on integer arrays")
+        if self.dtype in _complex_floating_dtypes:
+            raise TypeError("int is not allowed on complex floating-point arrays")
         res = self._array.__int__()
         return res
 
@@ -581,7 +602,7 @@ def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __le__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__le__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__le__")
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -603,7 +624,7 @@ def __lt__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __lt__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__lt__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__lt__")
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -626,7 +647,7 @@ def __mod__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __mod__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__mod__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__mod__")
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -808,7 +829,7 @@ def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __ifloordiv__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__ifloordiv__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__ifloordiv__")
         if other is NotImplemented:
             return other
         self._array.__ifloordiv__(other._array)
@@ -818,7 +839,7 @@ def __rfloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __rfloordiv__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__rfloordiv__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__rfloordiv__")
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -874,7 +895,7 @@ def __imod__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __imod__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__imod__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__imod__")
         if other is NotImplemented:
             return other
         self._array.__imod__(other._array)
@@ -884,7 +905,7 @@ def __rmod__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __rmod__.
         """
-        other = self._check_allowed_dtypes(other, "numeric", "__rmod__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__rmod__")
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)

array_api_strict/_dtypes.py

@@ -35,6 +35,7 @@
 _boolean_dtypes = (bool,)
 _real_floating_dtypes = (float32, float64)
 _floating_dtypes = (float32, float64, complex64, complex128)
+_complex_floating_dtypes = (complex64, complex128)
 _integer_dtypes = (int8, int16, int32, int64, uint8, uint16, uint32, uint64)
 _integer_or_boolean_dtypes = (
     bool,
@@ -47,6 +48,18 @@
     uint32,
     uint64,
 )
+_real_numeric_dtypes = (
+    float32,
+    float64,
+    int8,
+    int16,
+    int32,
+    int64,
+    uint8,
+    uint16,
+    uint32,
+    uint64,
+)
 _numeric_dtypes = (
     float32,
     float64,
@@ -64,6 +77,7 @@
 
 _dtype_categories = {
     "all": _all_dtypes,
+    "real numeric": _real_numeric_dtypes,
     "numeric": _numeric_dtypes,
     "integer": _integer_dtypes,
     "integer or boolean": _integer_or_boolean_dtypes,
@@ -144,7 +158,7 @@
     (complex64, complex64): complex64,
     (complex64, complex128): complex128,
     (complex128, complex64): complex128,
-    (complex128, complex64): complex128,
+    (complex128, complex128): complex128,
     (float32, complex64): complex64,
     (float32, complex128): complex128,
     (float64, complex64): complex128,

array_api_strict/tests/test_array_object.py

@@ -1,6 +1,6 @@
 import operator
 
-from numpy.testing import assert_raises
+from numpy.testing import assert_raises, suppress_warnings
 import numpy as np
 import pytest
 
@@ -9,9 +9,12 @@
 from numpy._dtypes import (
     _all_dtypes,
     _boolean_dtypes,
+    _real_floating_dtypes,
     _floating_dtypes,
+    _complex_floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
+    _real_numeric_dtypes,
     _numeric_dtypes,
     int8,
     int16,
@@ -85,13 +88,13 @@ def test_operators():
         "__add__": "numeric",
         "__and__": "integer_or_boolean",
         "__eq__": "all",
-        "__floordiv__": "numeric",
-        "__ge__": "numeric",
-        "__gt__": "numeric",
-        "__le__": "numeric",
+        "__floordiv__": "real numeric",
+        "__ge__": "real numeric",
+        "__gt__": "real numeric",
+        "__le__": "real numeric",
         "__lshift__": "integer",
-        "__lt__": "numeric",
-        "__mod__": "numeric",
+        "__lt__": "real numeric",
+        "__mod__": "real numeric",
         "__mul__": "numeric",
         "__ne__": "all",
         "__or__": "integer_or_boolean",
@@ -101,7 +104,6 @@ def test_operators():
         "__truediv__": "floating",
         "__xor__": "integer_or_boolean",
     }
-
     # Recompute each time because of in-place ops
     def _array_vals():
         for d in _integer_dtypes:
@@ -111,27 +113,28 @@ def _array_vals():
         for d in _floating_dtypes:
             yield asarray(1.0, dtype=d)
 
+
+    BIG_INT = int(1e30)
     for op, dtypes in binary_op_dtypes.items():
         ops = [op]
         if op not in ["__eq__", "__ne__", "__le__", "__ge__", "__lt__", "__gt__"]:
             rop = "__r" + op[2:]
             iop = "__i" + op[2:]
             ops += [rop, iop]
-        for s in [1, 1.0, False]:
+        for s in [1, 1.0, 1j, BIG_INT, False]:
             for _op in ops:
                 for a in _array_vals():
                     # Test array op scalar. From the spec, the following combinations
                     # are supported:
 
                     # - Python bool for a bool array dtype,
                     # - a Python int within the bounds of the given dtype for integer array dtypes,
-                    # - a Python int or float for floating-point array dtypes
-
-                    # We do not do bounds checking for int scalars, but rather use the default
-                    # NumPy behavior for casting in that case.
+                    # - a Python int or float for real floating-point array dtypes
+                    # - a Python int, float, or complex for complex floating-point array dtypes
 
                     if ((dtypes == "all"
                          or dtypes == "numeric" and a.dtype in _numeric_dtypes
+                         or dtypes == "real numeric" and a.dtype in _real_numeric_dtypes
                          or dtypes == "integer" and a.dtype in _integer_dtypes
                          or dtypes == "integer_or_boolean" and a.dtype in _integer_or_boolean_dtypes
                          or dtypes == "boolean" and a.dtype in _boolean_dtypes
@@ -141,10 +144,18 @@ def _array_vals():
                         # isinstance here.
                         and (a.dtype in _boolean_dtypes and type(s) == bool
                              or a.dtype in _integer_dtypes and type(s) == int
-                             or a.dtype in _floating_dtypes and type(s) in [float, int]
+                             or a.dtype in _real_floating_dtypes and type(s) in [float, int]
+                             or a.dtype in _complex_floating_dtypes and type(s) in [complex, float, int]
                         )):
-                        # Only test for no error
-                        getattr(a, _op)(s)
+                        if a.dtype in _integer_dtypes and s == BIG_INT:
+                            assert_raises(OverflowError, lambda: getattr(a, _op)(s))
+                        else:
+                            # Only test for no error
+                            with suppress_warnings() as sup:
+                                # ignore warnings from pow(BIG_INT)
+                                sup.filter(RuntimeWarning,
+                                           "invalid value encountered in power")
+                                getattr(a, _op)(s)
                     else:
                         assert_raises(TypeError, lambda: getattr(a, _op)(s))
 
@@ -174,8 +185,9 @@ def _array_vals():
                             # Ensure only those dtypes that are required for every operator are allowed.
                             elif (dtypes == "all" and (x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes
                                                       or x.dtype in _numeric_dtypes and y.dtype in _numeric_dtypes)
+                                or (dtypes == "real numeric" and x.dtype in _real_numeric_dtypes and y.dtype in _real_numeric_dtypes)
                                 or (dtypes == "numeric" and x.dtype in _numeric_dtypes and y.dtype in _numeric_dtypes)
-                                or dtypes == "integer" and x.dtype in _integer_dtypes and y.dtype in _numeric_dtypes
+                                or dtypes == "integer" and x.dtype in _integer_dtypes and y.dtype in _integer_dtypes
                                 or dtypes == "integer_or_boolean" and (x.dtype in _integer_dtypes and y.dtype in _integer_dtypes
                                                                        or x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes)
                                 or dtypes == "boolean" and x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes
@@ -263,31 +275,39 @@ def test_python_scalar_construtors():
     b = asarray(False)
     i = asarray(0)
     f = asarray(0.0)
+    c = asarray(0j)
 
     assert bool(b) == False
     assert int(i) == 0
     assert float(f) == 0.0
     assert operator.index(i) == 0
 
-    # bool/int/float should only be allowed on 0-D arrays.
+    # bool/int/float/complex should only be allowed on 0-D arrays.
     assert_raises(TypeError, lambda: bool(asarray([False])))
     assert_raises(TypeError, lambda: int(asarray([0])))
     assert_raises(TypeError, lambda: float(asarray([0.0])))
+    assert_raises(TypeError, lambda: complex(asarray([0j])))
     assert_raises(TypeError, lambda: operator.index(asarray([0])))
 
-    # bool/int/float should only be allowed on arrays of the corresponding
-    # dtype
-    assert_raises(ValueError, lambda: bool(i))
-    assert_raises(ValueError, lambda: bool(f))
+    # bool should work on all types of arrays
+    assert bool(b) is bool(i) is bool(f) is bool(c) is False
+
+    # int should fail on complex arrays
+    assert int(b) == int(i) == int(f) == 0
+    assert_raises(TypeError, lambda: int(c))
 
-    assert_raises(ValueError, lambda: int(b))
-    assert_raises(ValueError, lambda: int(f))
+    # float should fail on complex arrays
+    assert float(b) == float(i) == float(f) == 0.0
+    assert_raises(TypeError, lambda: float(c))
 
-    assert_raises(ValueError, lambda: float(b))
-    assert_raises(ValueError, lambda: float(i))
+    # complex should work on all types of arrays
+    assert complex(b) == complex(i) == complex(f) == complex(c) == 0j
 
+    # index should only work on integer arrays
+    assert operator.index(i) == 0
     assert_raises(TypeError, lambda: operator.index(b))
     assert_raises(TypeError, lambda: operator.index(f))
+    assert_raises(TypeError, lambda: operator.index(c))
 
 
 def test_device_property():