Improve scalar + array, and add complex support

Author: tadeu

numpy-stubs/__init__.pyi

@@ -42,12 +42,17 @@ class void:
 # and this is why we let int32 be a subclass of int64; and similarly for float32 and float64
 # the same logic applies when adding unsigned and signed values (uint + int -> int)
 
+class complex128(void, int):
+    def __complex__(self) -> complex: ...
+
+class complex64(complex128): ...
+
 # this would be the correct definition, but it makes `int` conflict with `float`
 # class float64(void, float): ...
-class float64(void, int):
+class float64(complex128):
     def __float__(self) -> float: ...
 
-class float32(float64): ...
+class float32(float64, complex64): ...
 class float16(float32): ...
 
 floating = float64
@@ -70,6 +75,8 @@ integer = int64
 _DType = TypeVar(
     "_DType",
     bool_,
+    complex64,
+    complex128,
     float16,
     float32,
     float64,
@@ -87,6 +94,8 @@ _DType = TypeVar(
 _DType2 = TypeVar(
     "_DType2",
     bool_,
+    complex64,
+    complex128,
     float16,
     float32,
     float64,
@@ -110,7 +119,7 @@ _ScalarLike = Union[_DType, str, int, float]
 _ConditionType = Union[ndarray[bool_], bool_, bool]
 newaxis: None = ...
 
-_AnyNum = Union[int, float, bool]
+_AnyNum = Union[int, float, bool, complex]
 # generic types that are only allowed to take on dtype values
 
 _Float = TypeVar("_Float", float16, float32, float64)
@@ -391,7 +400,7 @@ class ndarray(Generic[_DType]):
     @overload
     def __radd__(self, value: _DType) -> ndarray[_DType]: ...
     @overload
-    def __radd__(self, value: float) -> ndarray[_DType]: ...
+    def __radd__(self, value: float) -> ndarray[_DType2]: ...
     def __rand__(self, value: object) -> ndarray[_DType]: ...
     def __rdivmod__(self, value: object) -> Tuple[ndarray[_DType], ndarray[_DType]]: ...
     def __rfloordiv__(self, value: object) -> ndarray[_DType]: ...
@@ -506,6 +515,8 @@ def array(object: _NestedList[int]) -> ndarray[int64]: ...
 @overload
 def array(object: _NestedList[float]) -> ndarray[float64]: ...
 @overload
+def array(object: _NestedList[complex]) -> ndarray[complex64]: ...
+@overload
 def array(object: _NestedList[str]) -> ndarray[str_]: ...
 @overload
 def array(object: str) -> ndarray[str_]: ...

tests/numpy_test.py

@@ -8,6 +8,8 @@
 DType = TypeVar(
     "DType",
     np.bool_,
+    np.complex64,
+    np.complex128,
     np.float32,
     np.float64,
     np.int8,
@@ -306,6 +308,10 @@ def test_newaxis() -> None:
 
 
 def test_sum_scalar_before() -> None:
-    x = 273.15 + np.array([-0.1e2, -0.77e1])
+    x: np.ndarray[np.float64] = 273.15 + np.array([10, 20])
     assert isinstance(x, np.ndarray)
-    assert x.dtype == np.float64
+    assert_dtype(x, np.float64)
+
+    y: np.ndarray[np.complex128] = 10.0 + np.array([1j, 2j])
+    assert isinstance(y, np.ndarray)
+    assert_dtype(y, np.complex128)