MAINT: Fixed an issue with the return-dtype of ndarray.real and `imag

Bas van Beek · BvB93 · commit c66ffd98ea1c · 2021-06-20T15:56:42.000+02:00
The latter two would previously return complex arrays if the initial array was also complex
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -1633,6 +1633,14 @@ _ArrayTD64_co = NDArray[Union[bool_, integer[Any], timedelta64]]
 class _SupportsItem(Protocol[_T_co]):
     def item(self, __args: Any) -> _T_co: ...
 
+class _SupportsReal(Protocol[_T_co]):
+    @property
+    def real(self) -> _T_co: ...
+
+class _SupportsImag(Protocol[_T_co]):
+    @property
+    def imag(self) -> _T_co: ...
+
 class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
     @property
     def base(self) -> Optional[ndarray]: ...
@@ -1641,11 +1649,15 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
     @property
     def size(self) -> int: ...
     @property
-    def real(self: _ArraySelf) -> _ArraySelf: ...
+    def real(
+        self: NDArray[_SupportsReal[_ScalarType]],  # type: ignore[type-var]
+    ) -> ndarray[_ShapeType, dtype[_ScalarType]]: ...
     @real.setter
     def real(self, value: ArrayLike) -> None: ...
     @property
-    def imag(self: _ArraySelf) -> _ArraySelf: ...
+    def imag(
+        self: NDArray[_SupportsImag[_ScalarType]],  # type: ignore[type-var]
+    ) -> ndarray[_ShapeType, dtype[_ScalarType]]: ...
     @imag.setter
     def imag(self, value: ArrayLike) -> None: ...
     def __new__(
