Merge pull request #27753 from jorenham/typing/ndarray.real-imag

charris · web-flow · commit e0eee7ffc9d6 · 2024-11-14T10:57:17.000-07:00
TYP: Fix the annotations of ``ndarray.real`` and ``ndarray.imag``
diff --git a/numpy/__init__.pyi b/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"]
 
@@ -1823,35 +1826,39 @@ class _HasShapeAndDTypeWithItem(Protocol[_ShapeType_co, _T_co]):
     def dtype(self, /) -> _HashTypeWithItem[_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,
diff --git a/numpy/lib/_type_check_impl.pyi b/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]: ...
 
diff --git a/numpy/typing/tests/data/reveal/type_check.pyi b/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])
