Merge pull request #19286 from BvB93/type_check

ENH: Add annotations for `np.lib.type_check`

numpy/__init__.pyi

@@ -1636,6 +1636,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]: ...
@@ -1644,11 +1652,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__(

numpy/lib/type_check.pyi

@@ -1,19 +1,235 @@
-from typing import List
+import sys
+from typing import (
+    Any,
+    Container,
+    Iterable,
+    List,
+    overload,
+    Type,
+    TypeVar,
+)
+
+from numpy import (
+    dtype,
+    generic,
+    bool_,
+    floating,
+    float64,
+    complexfloating,
+    integer,
+)
+
+from numpy.typing import (
+    ArrayLike,
+    DTypeLike,
+    NBitBase,
+    NDArray,
+    _64Bit,
+    _SupportsDType,
+    _ScalarLike_co,
+    _NestedSequence,
+    _SupportsArray,
+    _DTypeLikeComplex,
+)
+
+if sys.version_info >= (3, 8):
+    from typing import Protocol, Literal as L
+else:
+    from typing_extensions import Protocol, Literal as L
+
+_T = TypeVar("_T")
+_T_co = TypeVar("_T_co", covariant=True)
+_SCT = TypeVar("_SCT", bound=generic)
+_NBit1 = TypeVar("_NBit1", bound=NBitBase)
+_NBit2 = TypeVar("_NBit2", bound=NBitBase)
+
+_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
+
+class _SupportsReal(Protocol[_T_co]):
+    @property
+    def real(self) -> _T_co: ...
+
+class _SupportsImag(Protocol[_T_co]):
+    @property
+    def imag(self) -> _T_co: ...
 
 __all__: List[str]
 
-def mintypecode(typechars, typeset=..., default=...): ...
-def asfarray(a, dtype = ...): ...
-def real(val): ...
-def imag(val): ...
-def iscomplex(x): ...
-def isreal(x): ...
-def iscomplexobj(x): ...
-def isrealobj(x): ...
-def nan_to_num(x, copy=..., nan=..., posinf=..., neginf=...): ...
-def real_if_close(a, tol=...): ...
-def typename(char): ...
-def common_type(*arrays): ...
-
-# NOTE: Deprecated
+def mintypecode(
+    typechars: Iterable[str | ArrayLike],
+    typeset: Container[str] = ...,
+    default: str = ...,
+) -> str: ...
+
+# `asfarray` ignores dtypes if they're not inexact
+
+@overload
+def asfarray(
+    a: object,
+    dtype: None | Type[float] = ...,
+) -> NDArray[float64]: ...
+@overload
+def asfarray(  # type: ignore[misc]
+    a: Any,
+    dtype: _DTypeLikeComplex,
+) -> NDArray[complexfloating[Any, Any]]: ...
+@overload
+def asfarray(
+    a: Any,
+    dtype: DTypeLike,
+) -> NDArray[floating[Any]]: ...
+
+@overload
+def real(val: _SupportsReal[_T]) -> _T: ...
+@overload
+def real(val: ArrayLike) -> NDArray[Any]: ...
+
+@overload
+def imag(val: _SupportsImag[_T]) -> _T: ...
+@overload
+def imag(val: ArrayLike) -> NDArray[Any]: ...
+
+@overload
+def iscomplex(x: _ScalarLike_co) -> bool_: ...  # type: ignore[misc]
+@overload
+def iscomplex(x: ArrayLike) -> NDArray[bool_]: ...
+
+@overload
+def isreal(x: _ScalarLike_co) -> bool_: ...  # type: ignore[misc]
+@overload
+def isreal(x: ArrayLike) -> NDArray[bool_]: ...
+
+def iscomplexobj(x: _SupportsDType[dtype[Any]] | ArrayLike) -> bool: ...
+
+def isrealobj(x: _SupportsDType[dtype[Any]] | ArrayLike) -> bool: ...
+
+@overload
+def nan_to_num(  # type: ignore[misc]
+    x: _SCT,
+    copy: bool = ...,
+    nan: float = ...,
+    posinf: None | float = ...,
+    neginf: None | float = ...,
+) -> _SCT: ...
+@overload
+def nan_to_num(
+    x: _ScalarLike_co,
+    copy: bool = ...,
+    nan: float = ...,
+    posinf: None | float = ...,
+    neginf: None | float = ...,
+) -> Any: ...
+@overload
+def nan_to_num(
+    x: _ArrayLike[_SCT],
+    copy: bool = ...,
+    nan: float = ...,
+    posinf: None | float = ...,
+    neginf: None | float = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def nan_to_num(
+    x: ArrayLike,
+    copy: bool = ...,
+    nan: float = ...,
+    posinf: None | float = ...,
+    neginf: None | float = ...,
+) -> NDArray[Any]: ...
+
+# If one passes a complex array to `real_if_close`, then one is reasonably
+# expected to verify the output dtype (so we can return an unsafe union here)
+
+@overload
+def real_if_close(  # type: ignore[misc]
+    a: _ArrayLike[complexfloating[_NBit1, _NBit1]],
+    tol: float = ...,
+) -> NDArray[floating[_NBit1]] | NDArray[complexfloating[_NBit1, _NBit1]]: ...
+@overload
+def real_if_close(
+    a: _ArrayLike[_SCT],
+    tol: float = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def real_if_close(
+    a: ArrayLike,
+    tol: float = ...,
+) -> NDArray[Any]: ...
+
+# NOTE: deprecated
 # def asscalar(a): ...
+
+@overload
+def typename(char: L['S1']) -> L['character']: ...
+@overload
+def typename(char: L['?']) -> L['bool']: ...
+@overload
+def typename(char: L['b']) -> L['signed char']: ...
+@overload
+def typename(char: L['B']) -> L['unsigned char']: ...
+@overload
+def typename(char: L['h']) -> L['short']: ...
+@overload
+def typename(char: L['H']) -> L['unsigned short']: ...
+@overload
+def typename(char: L['i']) -> L['integer']: ...
+@overload
+def typename(char: L['I']) -> L['unsigned integer']: ...
+@overload
+def typename(char: L['l']) -> L['long integer']: ...
+@overload
+def typename(char: L['L']) -> L['unsigned long integer']: ...
+@overload
+def typename(char: L['q']) -> L['long long integer']: ...
+@overload
+def typename(char: L['Q']) -> L['unsigned long long integer']: ...
+@overload
+def typename(char: L['f']) -> L['single precision']: ...
+@overload
+def typename(char: L['d']) -> L['double precision']: ...
+@overload
+def typename(char: L['g']) -> L['long precision']: ...
+@overload
+def typename(char: L['F']) -> L['complex single precision']: ...
+@overload
+def typename(char: L['D']) -> L['complex double precision']: ...
+@overload
+def typename(char: L['G']) -> L['complex long double precision']: ...
+@overload
+def typename(char: L['S']) -> L['string']: ...
+@overload
+def typename(char: L['U']) -> L['unicode']: ...
+@overload
+def typename(char: L['V']) -> L['void']: ...
+@overload
+def typename(char: L['O']) -> L['object']: ...
+
+@overload
+def common_type(  # type: ignore[misc]
+    *arrays: _SupportsDType[dtype[
+        integer[Any]
+    ]]
+) -> Type[floating[_64Bit]]: ...
+@overload
+def common_type(  # type: ignore[misc]
+    *arrays: _SupportsDType[dtype[
+        floating[_NBit1]
+    ]]
+) -> Type[floating[_NBit1]]: ...
+@overload
+def common_type(  # type: ignore[misc]
+    *arrays: _SupportsDType[dtype[
+        integer[Any] | floating[_NBit1]
+    ]]
+) -> Type[floating[_NBit1 | _64Bit]]: ...
+@overload
+def common_type(  # type: ignore[misc]
+    *arrays: _SupportsDType[dtype[
+        floating[_NBit1] | complexfloating[_NBit2, _NBit2]
+    ]]
+) -> Type[complexfloating[_NBit1 | _NBit2, _NBit1 | _NBit2]]: ...
+@overload
+def common_type(
+    *arrays: _SupportsDType[dtype[
+        integer[Any] | floating[_NBit1] | complexfloating[_NBit2, _NBit2]
+    ]]
+) -> Type[complexfloating[_64Bit | _NBit1 | _NBit2, _64Bit | _NBit1 | _NBit2]]: ...

numpy/typing/tests/data/fail/type_check.py

@@ -0,0 +1,13 @@
+import numpy as np
+import numpy.typing as npt
+
+DTYPE_i8: np.dtype[np.int64]
+
+np.mintypecode(DTYPE_i8)  # E: incompatible type
+np.iscomplexobj(DTYPE_i8)  # E: incompatible type
+np.isrealobj(DTYPE_i8)  # E: incompatible type
+
+np.typename(DTYPE_i8)  # E: No overload variant
+np.typename("invalid")  # E: No overload variant
+
+np.common_type(np.timedelta64())  # E: incompatible type

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

@@ -0,0 +1,73 @@
+from typing import List
+import numpy as np
+import numpy.typing as npt
+
+f8: np.float64
+f: float
+
+# NOTE: Avoid importing the platform specific `np.float128` type
+AR_i8: npt.NDArray[np.int64]
+AR_i4: npt.NDArray[np.int32]
+AR_f2: npt.NDArray[np.float16]
+AR_f8: npt.NDArray[np.float64]
+AR_f16: npt.NDArray[np.floating[npt._128Bit]]
+AR_c8: npt.NDArray[np.complex64]
+AR_c16: npt.NDArray[np.complex128]
+
+AR_LIKE_f: List[float]
+
+class RealObj:
+    real: slice
+
+class ImagObj:
+    imag: slice
+
+reveal_type(np.mintypecode(["f8"], typeset="qfQF"))
+
+reveal_type(np.asfarray(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.asfarray(AR_LIKE_f))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.asfarray(AR_f8, dtype="c16"))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+reveal_type(np.asfarray(AR_f8, dtype="i8"))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+
+reveal_type(np.real(RealObj()))  # E: slice
+reveal_type(np.real(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.real(AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.real(AR_LIKE_f))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.imag(ImagObj()))  # E: slice
+reveal_type(np.imag(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.imag(AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.imag(AR_LIKE_f))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.iscomplex(f8))  # E: numpy.bool_
+reveal_type(np.iscomplex(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+reveal_type(np.iscomplex(AR_LIKE_f))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+
+reveal_type(np.isreal(f8))  # E: numpy.bool_
+reveal_type(np.isreal(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+reveal_type(np.isreal(AR_LIKE_f))  # E: numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+
+reveal_type(np.iscomplexobj(f8))  # E: bool
+reveal_type(np.isrealobj(f8))  # E: bool
+
+reveal_type(np.nan_to_num(f8))  # E: {float64}
+reveal_type(np.nan_to_num(f, copy=True))  # E: Any
+reveal_type(np.nan_to_num(AR_f8, nan=1.5))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.nan_to_num(AR_LIKE_f, posinf=9999))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.real_if_close(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.real_if_close(AR_c16))  # E: Union[numpy.ndarray[Any, numpy.dtype[{float64}]], numpy.ndarray[Any, numpy.dtype[{complex128}]]]
+reveal_type(np.real_if_close(AR_c8))  # E: Union[numpy.ndarray[Any, numpy.dtype[{float32}]], numpy.ndarray[Any, numpy.dtype[{complex64}]]]
+reveal_type(np.real_if_close(AR_LIKE_f))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.typename("h"))  # E: Literal['short']
+reveal_type(np.typename("B"))  # E: Literal['unsigned char']
+reveal_type(np.typename("V"))  # E: Literal['void']
+reveal_type(np.typename("S1"))  # E: Literal['character']
+
+reveal_type(np.common_type(AR_i4))  # E: Type[{float64}]
+reveal_type(np.common_type(AR_f2))  # E: Type[{float16}]
+reveal_type(np.common_type(AR_f2, AR_i4))  # E: Type[{float64}]
+reveal_type(np.common_type(AR_f16, AR_i4))  # E: Type[{float128}]
+reveal_type(np.common_type(AR_c8, AR_f2))  # E: Type[{complex64}]
+reveal_type(np.common_type(AR_f2, AR_c8, AR_i4))  # E: Type[{complex128}]