TYP: Fix ndarray.real and .imag

Author: jorenham

numpy/__init__.pyi

@@ -1782,6 +1782,9 @@ else:
 _T = TypeVar("_T")
 _T_co = TypeVar("_T_co", covariant=True)
 _T_contra = TypeVar("_T_contra", contravariant=True)
+_RealT_co = TypeVar("_RealT_co", covariant=True)
+_ImagT_co = TypeVar("_ImagT_co", covariant=True)
+
 _2Tuple: TypeAlias = tuple[_T, _T]
 _CastingKind: TypeAlias = L["no", "equiv", "safe", "same_kind", "unsafe"]
 
@@ -1802,35 +1805,39 @@ class _SupportsItem(Protocol[_T_co]):
     def item(self, args: Any, /) -> _T_co: ...
 
 @type_check_only
-class _SupportsReal(Protocol[_T_co]):
+class _HasRealAndImag(Protocol[_RealT_co, _ImagT_co]):
+    @property
+    def real(self, /) -> _RealT_co: ...
     @property
-    def real(self) -> _T_co: ...
+    def imag(self, /) -> _ImagT_co: ...
 
 @type_check_only
-class _SupportsImag(Protocol[_T_co]):
+class _HasTypeWithRealAndImag(Protocol[_RealT_co, _ImagT_co]):
     @property
-    def imag(self) -> _T_co: ...
+    def type(self, /) -> type[_HasRealAndImag[_RealT_co, _ImagT_co]]: ...
+
+@type_check_only
+class _HasDTypeWithRealAndImag(Protocol[_RealT_co, _ImagT_co]):
+    @property
+    def dtype(self, /) -> _HasTypeWithRealAndImag[_RealT_co, _ImagT_co]: ...
 
 class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
-    __hash__: ClassVar[None]
+    __hash__: ClassVar[None]  # type: ignore[assignment]  # pyright: ignore[reportIncompatibleMethodOverride]
     @property
     def base(self) -> None | NDArray[Any]: ...
     @property
     def ndim(self) -> int: ...
     @property
     def size(self) -> int: ...
     @property
-    def real(
-        self: ndarray[_ShapeType_co, dtype[_SupportsReal[_ScalarType]]],  # type: ignore[type-var]
-    ) -> ndarray[_ShapeType_co, _dtype[_ScalarType]]: ...
+    def real(self: _HasDTypeWithRealAndImag[_SCT, object], /) -> ndarray[_ShapeType_co, dtype[_SCT]]: ...
     @real.setter
-    def real(self, value: ArrayLike) -> None: ...
+    def real(self, value: ArrayLike, /) -> None: ...
     @property
-    def imag(
-        self: ndarray[_ShapeType_co, dtype[_SupportsImag[_ScalarType]]],  # type: ignore[type-var]
-    ) -> ndarray[_ShapeType_co, _dtype[_ScalarType]]: ...
+    def imag(self: _HasDTypeWithRealAndImag[object, _SCT], /) -> ndarray[_ShapeType_co, dtype[_SCT]]: ...
     @imag.setter
-    def imag(self, value: ArrayLike) -> None: ...
+    def imag(self, value: ArrayLike, /) -> None: ...
+
     def __new__(
         cls,
         shape: _ShapeLike,

numpy/lib/_type_check_impl.pyi

@@ -3,8 +3,7 @@ from typing import Literal as L, Any, overload, TypeVar
 
 import numpy as np
 from numpy import (
-    _SupportsImag,
-    _SupportsReal,
+    _HasRealAndImag,
     dtype,
     generic,
     floating,
@@ -50,12 +49,12 @@ def mintypecode(
 ) -> str: ...
 
 @overload
-def real(val: _SupportsReal[_T]) -> _T: ...
+def real(val: _HasRealAndImag[_T, Any]) -> _T: ...
 @overload
 def real(val: ArrayLike) -> NDArray[Any]: ...
 
 @overload
-def imag(val: _SupportsImag[_T]) -> _T: ...
+def imag(val: _HasRealAndImag[Any, _T]) -> _T: ...
 @overload
 def imag(val: ArrayLike) -> NDArray[Any]: ...
 

numpy/typing/tests/data/reveal/type_check.pyi

@@ -20,20 +20,18 @@ AR_c16: npt.NDArray[np.complex128]
 
 AR_LIKE_f: list[float]
 
-class RealObj:
+class ComplexObj:
     real: slice
-
-class ImagObj:
     imag: slice
 
 assert_type(np.mintypecode(["f8"], typeset="qfQF"), str)
 
-assert_type(np.real(RealObj()), slice)
+assert_type(np.real(ComplexObj()), slice)
 assert_type(np.real(AR_f8), npt.NDArray[np.float64])
 assert_type(np.real(AR_c16), npt.NDArray[np.float64])
 assert_type(np.real(AR_LIKE_f), npt.NDArray[Any])
 
-assert_type(np.imag(ImagObj()), slice)
+assert_type(np.imag(ComplexObj()), slice)
 assert_type(np.imag(AR_f8), npt.NDArray[np.float64])
 assert_type(np.imag(AR_c16), npt.NDArray[np.float64])
 assert_type(np.imag(AR_LIKE_f), npt.NDArray[Any])