Merge pull request #320 from guan404ming/1-to-1-bool

🏷️ ufunc annotations for `is{finite,inf,nan.nat}`

Author: jorenham

src/numpy-stubs/__init__.pyi

@@ -307,7 +307,7 @@ from ._typing._char_codes import (
     _UnsignedIntegerCodes,
     _VoidCodes,
 )
-from ._typing._ufunc import _Call21Bool, _gufunc_2_1, _ufunc_1_1, _ufunc_1_2, _ufunc_2_1, _ufunc_2_2
+from ._typing._ufunc import _Call11Bool, _Call21Bool, _gufunc_2_1, _ufunc_1_1, _ufunc_1_2, _ufunc_2_1, _ufunc_2_2
 from .lib import scimath as emath
 from .lib._arraypad_impl import pad
 from .lib._arraysetops_impl import (
@@ -6989,10 +6989,10 @@ expm1: Final[_ufunc_1_1] = ...
 fabs: Final[_ufunc_1_1] = ...
 floor: Final[_ufunc_1_1] = ...
 invert: Final[_ufunc_1_1] = ...
-isfinite: Final[_ufunc_1_1] = ...
-isinf: Final[_ufunc_1_1] = ...
-isnan: Final[_ufunc_1_1] = ...
-isnat: Final[_ufunc_1_1] = ...
+isfinite: Final[_ufunc_1_1[_Call11Bool]] = ...
+isinf: Final[_ufunc_1_1[_Call11Bool]] = ...
+isnan: Final[_ufunc_1_1[_Call11Bool]] = ...
+isnat: Final[_ufunc_1_1[_Call11Bool]] = ...
 log: Final[_ufunc_1_1] = ...
 log2: Final[_ufunc_1_1] = ...
 log10: Final[_ufunc_1_1] = ...

src/numpy-stubs/_typing/_ufunc.pyi

@@ -14,7 +14,7 @@ from typing_extensions import Never, TypeAliasType, TypeVar, Unpack
 
 import numpy as np
 from numpy import _CastingKind, _OrderKACF  # noqa: ICN003
-from numpy._typing import _NestedSequence
+from numpy._typing import _DTypeLikeBool, _NestedSequence
 
 from ._array_like import ArrayLike, NDArray, _ArrayLikeBool_co, _ArrayLikeInt_co
 from ._dtype_like import DTypeLike, _DTypeLike
@@ -178,6 +178,49 @@ class _Call11(Protocol):
         **kwds: Unpack[_Kwargs2],
     ) -> Any: ...
 
+@type_check_only
+class _Call11Bool(Protocol):
+    @overload  # (scalar) -> bool
+    def __call__(
+        self,
+        x: _ScalarLike_co,
+        /,
+        out: None = None,
+        *,
+        dtype: _DTypeLikeBool | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> np.bool: ...
+    @overload  # (array-like, out: T) -> T
+    def __call__(
+        self,
+        x: ArrayLike,
+        /,
+        out: _Out1[_ArrayT],
+        *,
+        dtype: _DTypeLikeBool | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> _ArrayT: ...
+    @overload  # (array) -> Array[bool] | bool
+    def __call__(
+        self,
+        x: _AnyArray,
+        /,
+        out: _Out1[_AnyArray] | None = None,
+        *,
+        dtype: _DTypeLikeBool | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> NDArray[np.bool]: ...
+    @overload  # (array-like) -> Array[bool] | bool
+    def __call__(
+        self,
+        x: ArrayLike,
+        /,
+        out: _Out1[_AnyArray] | None = None,
+        *,
+        dtype: _DTypeLikeBool | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> NDArray[np.bool] | np.bool: ...
+
 @type_check_only
 class _Call12(Protocol):
     @overload

test/static/accept/ufuncs.pyi

@@ -8,6 +8,7 @@ i8: np.int64
 f8: np.float64
 AR_f8: npt.NDArray[np.float64]
 AR_i8: npt.NDArray[np.int64]
+AR_bool: npt.NDArray[np.bool_]
 
 assert_type(np.absolute.types, list[str])
 
@@ -36,3 +37,19 @@ assert_type(np.vecdot(AR_f8, AR_f8), Any)
 
 assert_type(np.bitwise_count(i8), Any)
 assert_type(np.bitwise_count(AR_i8), npt.NDArray[Any])
+
+assert_type(np.isnan(f8), np.bool_)
+assert_type(np.isnan(AR_f8), npt.NDArray[np.bool_])
+assert_type(np.isnan(AR_f8, out=AR_bool), npt.NDArray[np.bool_])
+
+assert_type(np.isnat(f8), np.bool_)
+assert_type(np.isnat(AR_f8), npt.NDArray[np.bool_])
+assert_type(np.isnat(AR_f8, out=AR_bool), npt.NDArray[np.bool_])
+
+assert_type(np.isinf(f8), np.bool_)
+assert_type(np.isinf(AR_f8), npt.NDArray[np.bool_])
+assert_type(np.isinf(AR_f8, out=AR_bool), npt.NDArray[np.bool_])
+
+assert_type(np.isfinite(f8), np.bool_)
+assert_type(np.isfinite(AR_f8), npt.NDArray[np.bool_])
+assert_type(np.isfinite(AR_f8, out=AR_bool), npt.NDArray[np.bool_])

test/static/reject/ufuncs.pyi

@@ -1,6 +1,7 @@
 import numpy as np
 import numpy.typing as npt
 
+i8: np.int64
 AR_f8: npt.NDArray[np.float64]
 
 np.sin.nin + "foo"  # type: ignore[operator]  # pyright: ignore[reportOperatorIssue]
@@ -40,3 +41,8 @@ np.matmul.accumulate()  # type: ignore[call-arg]  # pyright: ignore[reportCallIs
 np.frexp.at()  # type: ignore[call-arg, misc]  # pyright: ignore[reportCallIssue]
 np.divmod.at()  # type: ignore[call-arg, misc]  # pyright: ignore[reportCallIssue]
 np.matmul.at()  # type: ignore[call-arg, misc]  # pyright: ignore[reportCallIssue]
+
+np.isnan(i8, dtype=int)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType, reportCallIssue]
+np.isnat(i8, dtype=np.int64)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType, reportCallIssue]
+np.isinf(i8, dtype=np.int64)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType, reportCallIssue]
+np.isfinite(i8, dtype=np.int64)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType, reportCallIssue]