ENH: typing: Make np.complexfloating generic w.r.t. np.floating (#17172)

BvB93 · web-flow · commit 3fbc84a5662f · 2020-09-01T13:10:36.000+01:00
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -462,6 +462,8 @@ class uint64(unsignedinteger):
 class inexact(number): ...  # type: ignore
 class floating(inexact, _real_generic): ...  # type: ignore
 
+_FloatType = TypeVar('_FloatType', bound=floating)
+
 class float16(floating):
     def __init__(self, __value: Optional[SupportsFloat] = ...) -> None: ...
 
@@ -471,27 +473,24 @@ class float32(floating):
 class float64(floating):
     def __init__(self, __value: Optional[SupportsFloat] = ...) -> None: ...
 
-class complexfloating(inexact): ...  # type: ignore
+class complexfloating(inexact, Generic[_FloatType]):  # type: ignore
+    @property
+    def real(self) -> _FloatType: ...
+    @property
+    def imag(self) -> _FloatType: ...
+    def __abs__(self) -> _FloatType: ...  # type: ignore[override]
 
-class complex64(complexfloating):
+class complex64(complexfloating[float32]):
     def __init__(
         self,
         __value: Union[None, SupportsInt, SupportsFloat, SupportsComplex] = ...
     ) -> None: ...
-    @property
-    def real(self) -> float32: ...
-    @property
-    def imag(self) -> float32: ...
 
-class complex128(complexfloating):
+class complex128(complexfloating[float64]):
     def __init__(
         self,
         __value: Union[None, SupportsInt, SupportsFloat, SupportsComplex] = ...
     ) -> None: ...
-    @property
-    def real(self) -> float64: ...
-    @property
-    def imag(self) -> float64: ...
 
 class flexible(_real_generic): ...  # type: ignore
 
diff --git a/numpy/tests/typing/reveal/scalars.py b/numpy/tests/typing/reveal/scalars.py
@@ -28,3 +28,6 @@
 reveal_type(td / 1.0)  # E: numpy.timedelta64
 reveal_type(td / td)  # E: float
 reveal_type(td % td)  # E: numpy.timedelta64
+
+reveal_type(np.complex64().real)  # E: numpy.float32
+reveal_type(np.complex128().imag)  # E: numpy.float64
