TYP: Concrete complex128 scalar type with builtins.complex as a base class

Author: jorenham

numpy/__init__.pyi

@@ -2845,6 +2845,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     # operand. An exception to this rule are unsigned integers though, which
     # also accepts a signed integer for the right operand as long it is a 0D
     # object and its value is >= 0
+    # NOTE: Due to a mypy bug, overloading on e.g. `self: NDArray[SCT_floating]` won't
+    # work, as this will lead to `false negatives` when using these inplace ops.
     @overload
     def __iadd__(self: NDArray[_UnknownType], other: _ArrayLikeUnknown, /) -> NDArray[Any]: ...
     @overload
@@ -2858,6 +2860,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     @overload
     def __iadd__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co, /) -> NDArray[floating[_NBit1]]: ...
     @overload
+    def __iadd__(self: NDArray[complex128], other: _ArrayLikeComplex_co, /) -> NDArray[complex128]: ...
+    @overload
     def __iadd__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co, /) -> NDArray[complexfloating[_NBit1, _NBit1]]: ...
     @overload
     def __iadd__(self: NDArray[timedelta64], other: _ArrayLikeTD64_co, /) -> NDArray[timedelta64]: ...
@@ -2877,6 +2881,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     @overload
     def __isub__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co, /) -> NDArray[floating[_NBit1]]: ...
     @overload
+    def __isub__(self: NDArray[complex128], other: _ArrayLikeComplex_co, /) -> NDArray[complex128]: ...
+    @overload
     def __isub__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co, /) -> NDArray[complexfloating[_NBit1, _NBit1]]: ...
     @overload
     def __isub__(self: NDArray[timedelta64], other: _ArrayLikeTD64_co, /) -> NDArray[timedelta64]: ...
@@ -2898,6 +2904,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     @overload
     def __imul__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co, /) -> NDArray[floating[_NBit1]]: ...
     @overload
+    def __imul__(self: NDArray[complex128], other: _ArrayLikeComplex_co, /) -> NDArray[complex128]: ...
+    @overload
     def __imul__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co, /) -> NDArray[complexfloating[_NBit1, _NBit1]]: ...
     @overload
     def __imul__(self: NDArray[timedelta64], other: _ArrayLikeFloat_co, /) -> NDArray[timedelta64]: ...
@@ -2911,6 +2919,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     @overload
     def __itruediv__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co, /) -> NDArray[floating[_NBit1]]: ...
     @overload
+    def __itruediv__(self: NDArray[complex128], other: _ArrayLikeComplex_co, /) -> NDArray[complex128]: ...
+    @overload
     def __itruediv__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co, /) -> NDArray[complexfloating[_NBit1, _NBit1]]: ...
     @overload
     def __itruediv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co, /) -> NoReturn: ...
@@ -2930,6 +2940,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     @overload
     def __ifloordiv__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co, /) -> NDArray[floating[_NBit1]]: ...
     @overload
+    def __ifloordiv__(self: NDArray[complex128], other: _ArrayLikeComplex_co, /) -> NDArray[complex128]: ...
+    @overload
     def __ifloordiv__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co, /) -> NDArray[complexfloating[_NBit1, _NBit1]]: ...
     @overload
     def __ifloordiv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co, /) -> NoReturn: ...
@@ -2949,6 +2961,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     @overload
     def __ipow__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co, /) -> NDArray[floating[_NBit1]]: ...
     @overload
+    def __ipow__(self: NDArray[complex128], other: _ArrayLikeComplex_co, /) -> NDArray[complex128]: ...
+    @overload
     def __ipow__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co, /) -> NDArray[complexfloating[_NBit1, _NBit1]]: ...
     @overload
     def __ipow__(self: NDArray[object_], other: Any, /) -> NDArray[object_]: ...
@@ -3032,6 +3046,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     @overload
     def __imatmul__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co, /) -> NDArray[floating[_NBit1]]: ...
     @overload
+    def __imatmul__(self: NDArray[complex128], other: _ArrayLikeComplex_co, /) -> NDArray[complex128]: ...
+    @overload
     def __imatmul__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co, /) -> NDArray[complexfloating[_NBit1, _NBit1]]: ...
     @overload
     def __imatmul__(self: NDArray[object_], other: Any, /) -> NDArray[object_]: ...
@@ -3534,8 +3550,7 @@ uint: TypeAlias = uintp
 ulong: TypeAlias = unsignedinteger[_NBitLong]
 ulonglong: TypeAlias = unsignedinteger[_NBitLongLong]
 
-class inexact(number[_NBit1]):  # type: ignore
-    def __getnewargs__(self: inexact[_64Bit]) -> tuple[float, ...]: ...
+class inexact(number[_NBit1]): ...  # type: ignore[misc]
 
 _IntType = TypeVar("_IntType", bound=integer[Any])
 
@@ -3571,6 +3586,7 @@ float32: TypeAlias = floating[_32Bit]
 
 # NOTE: `_64Bit` is equivalent to `_64Bit | _32Bit | _16Bit | _8Bit`
 _Float64_co: TypeAlias = float | floating[_64Bit] | integer[_64Bit] | np.bool
+_Complex128_co: TypeAlias = complex | complexfloating[_64Bit, _64Bit] | _Float64_co
 
 # either a C `double`, `float`, or `longdouble`
 class float64(floating[_64Bit], float):  # type: ignore[misc]
@@ -3588,77 +3604,101 @@ class float64(floating[_64Bit], float):  # type: ignore[misc]
     @overload
     def __add__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __add__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __add__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __add__(self, other: complex, /) -> float64 | complex128: ...
     @overload
     def __radd__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __radd__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __radd__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __radd__(self, other: complex, /) -> float64 | complex128: ...
 
     @overload
     def __sub__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __sub__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __sub__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __sub__(self, other: complex, /) -> float64 | complex128: ...
     @overload
     def __rsub__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __rsub__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __rsub__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __rsub__(self, other: complex, /) -> float64 | complex128: ...
 
     @overload
     def __mul__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __mul__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __mul__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __mul__(self, other: complex, /) -> float64 | complex128: ...
     @overload
     def __rmul__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __rmul__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __rmul__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __rmul__(self, other: complex, /) -> float64 | complex128: ...
 
     @overload
     def __truediv__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __truediv__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __truediv__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __truediv__(self, other: complex, /) -> float64 | complex128: ...
     @overload
     def __rtruediv__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __rtruediv__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __rtruediv__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __rtruediv__(self, other: complex, /) -> float64 | complex128: ...
 
     @overload
     def __floordiv__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __floordiv__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __floordiv__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __floordiv__(self, other: complex, /) -> float64 | complex128: ...
     @overload
     def __rfloordiv__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __rfloordiv__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __rfloordiv__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __rfloordiv__(self, other: complex, /) -> float64 | complex128: ...
 
     @overload
     def __pow__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __pow__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __pow__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __pow__(self, other: complex, /) -> float64 | complex128: ...
     @overload
     def __rpow__(self, other: _Float64_co, /) -> float64: ...
     @overload
+    def __rpow__(self, other: complexfloating[_64Bit, _64Bit], /) -> complex128: ...
+    @overload
     def __rpow__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
     @overload
     def __rpow__(self, other: complex, /) -> float64 | complex128: ...
@@ -3681,16 +3721,13 @@ longdouble: TypeAlias = floating[_NBitLongDouble]
 
 class complexfloating(inexact[_NBit1], Generic[_NBit1, _NBit2]):
     def __init__(self, value: _ComplexValue = ..., /) -> None: ...
-    def item(
-        self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /,
-    ) -> complex: ...
+    def item(self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /) -> complex: ...
     def tolist(self) -> complex: ...
     @property
     def real(self) -> floating[_NBit1]: ...  # type: ignore[override]
     @property
     def imag(self) -> floating[_NBit2]: ...  # type: ignore[override]
     def __abs__(self) -> floating[_NBit1]: ...  # type: ignore[override]
-    def __getnewargs__(self: complex128) -> tuple[float, float]: ...
     # NOTE: Deprecated
     # def __round__(self, ndigits=...): ...
     __add__: _ComplexOp[_NBit1]
@@ -3705,7 +3742,49 @@ class complexfloating(inexact[_NBit1], Generic[_NBit1, _NBit2]):
     __rpow__: _ComplexOp[_NBit1]
 
 complex64: TypeAlias = complexfloating[_32Bit, _32Bit]
-complex128: TypeAlias = complexfloating[_64Bit, _64Bit]
+
+class complex128(complexfloating[_64Bit, _64Bit], complex):
+    def __getnewargs__(self, /) -> tuple[float, float]: ...
+
+    # overrides for `floating` and `builtins.float` compatibility
+    @property
+    def real(self) -> float64: ...
+    @property
+    def imag(self) -> float64: ...
+    def __abs__(self) -> float64: ...
+    def conjugate(self) -> Self: ...
+
+    # complex128-specific operator overrides
+    @overload
+    def __add__(self, other: _Complex128_co, /) -> complex128: ...
+    @overload
+    def __add__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
+    def __radd__(self, other: _Complex128_co, /) -> complex128: ...
+
+    @overload
+    def __sub__(self, other: _Complex128_co, /) -> complex128: ...
+    @overload
+    def __sub__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
+    def __rsub__(self, other: _Complex128_co, /) -> complex128: ...
+
+    @overload
+    def __mul__(self, other: _Complex128_co, /) -> complex128: ...
+    @overload
+    def __mul__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
+    def __rmul__(self, other: _Complex128_co, /) -> complex128: ...
+
+    @overload
+    def __truediv__(self, other: _Complex128_co, /) -> complex128: ...
+    @overload
+    def __truediv__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
+    def __rtruediv__(self, other: _Complex128_co, /) -> complex128: ...
+
+    @overload
+    def __pow__(self, other: _Complex128_co, /) -> complex128: ...
+    @overload
+    def __pow__(self, other: complexfloating[_NBit1, _NBit2], /) -> complexfloating[_NBit1 | _64Bit, _NBit2 | _64Bit]: ...
+    def __rpow__(self, other: _Complex128_co, /) -> complex128: ...
+
 
 csingle: TypeAlias = complexfloating[_NBitSingle, _NBitSingle]
 cdouble: TypeAlias = complexfloating[_NBitDouble, _NBitDouble]

numpy/typing/tests/data/pass/arithmetic.py

@@ -2,6 +2,7 @@
 
 from typing import Any
 import numpy as np
+import numpy.typing as npt
 import pytest
 
 c16 = np.complex128(1)
@@ -57,14 +58,14 @@ def __rpow__(self, value: Any) -> Object:
         return self
 
 
-AR_b: np.ndarray[Any, np.dtype[np.bool]] = np.array([True])
-AR_u: np.ndarray[Any, np.dtype[np.uint32]] = np.array([1], dtype=np.uint32)
-AR_i: np.ndarray[Any, np.dtype[np.int64]] = np.array([1])
-AR_f: np.ndarray[Any, np.dtype[np.float64]] = np.array([1.0])
-AR_c: np.ndarray[Any, np.dtype[np.complex128]] = np.array([1j])
-AR_m: np.ndarray[Any, np.dtype[np.timedelta64]] = np.array([np.timedelta64(1, "D")])
-AR_M: np.ndarray[Any, np.dtype[np.datetime64]] = np.array([np.datetime64(1, "D")])
-AR_O: np.ndarray[Any, np.dtype[np.object_]] = np.array([Object()])
+AR_b: npt.NDArray[np.bool] = np.array([True])
+AR_u: npt.NDArray[np.uint32] = np.array([1], dtype=np.uint32)
+AR_i: npt.NDArray[np.int64] = np.array([1])
+AR_f: npt.NDArray[np.float64] = np.array([1.0])
+AR_c: npt.NDArray[np.complex128] = np.array([1j])
+AR_m: npt.NDArray[np.timedelta64] = np.array([np.timedelta64(1, "D")])
+AR_M: npt.NDArray[np.datetime64] = np.array([np.datetime64(1, "D")])
+AR_O: npt.NDArray[np.object_] = np.array([Object()])
 
 AR_LIKE_b = [True]
 AR_LIKE_u = [np.uint32(1)]