🏷️ ufunc annotations for spacing (#363)

* 🏷️ stub `umath.spacing`

* ✅ add test for `umath.spacing`

* 🏷️ enhance annatation with float64

Author: guan404ming

src/numpy-stubs/_core/umath.pyi

@@ -21,11 +21,14 @@ from numpy._typing import (
     NDArray,
     _ArrayLike,
     _ArrayLikeBool_co,
+    _ArrayLikeFloat_co,
     _ArrayLikeInt_co,
     _ArrayLikeNumber_co,
     _ArrayLikeObject_co,
     _DTypeLike,
     _DTypeLikeBool,
+    _DTypeLikeFloat,
+    _FloatLike_co,
     _NestedSequence,
     _NumberLike_co,
     _ScalarLike_co,
@@ -357,6 +360,59 @@ class _Call11Bool(Protocol):
         **kwds: Unpack[_Kwargs2],
     ) -> NDArray[np.bool] | np.bool: ...
 
+@type_check_only
+class _Call11Float(Protocol):
+    @overload  # (float) -> float64
+    def __call__(
+        self,
+        x: float,
+        /,
+        out: None = None,
+        *,
+        dtype: _DTypeLikeFloat | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> np.float64: ...
+    @overload  # (scalar) -> float
+    def __call__(
+        self,
+        x: _FloatLike_co,
+        /,
+        out: None = None,
+        *,
+        dtype: _DTypeLikeFloat | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> np.floating: ...
+    @overload  # (array-like, out: T) -> T
+    def __call__(
+        self,
+        x: _ArrayLikeFloat_co,
+        /,
+        out: _Out1[_ArrayT],
+        *,
+        dtype: _DTypeLikeFloat | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> _ArrayT: ...
+    @overload  # (NDArray[float64] | _NestedSequence[float]) -> NDArray[float64]
+    def __call__(
+        self,
+        x: NDArray[np.float64] | _NestedSequence[float],
+        /,
+        out: _Out1[NDArray[np.float64]] | None = None,
+        *,
+        dtype: _DTypeLikeFloat | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> NDArray[np.float64]: ...
+    @overload  # (array) -> Array[float]
+    def __call__(
+        self,
+        x: NDArray[np.floating] | _NestedSequence[np.floating],
+        /,
+        out: _Out1[NDArray[np.floating]] | None = None,
+        *,
+        dtype: _DTypeLikeFloat | None = None,
+        **kwds: Unpack[_Kwargs2],
+    ) -> NDArray[np.floating]: ...
+
 @type_check_only
 class _Call11Isnat(Protocol):
     @overload  # (scalar) -> bool
@@ -1151,8 +1207,8 @@ class _ReduceAt2(Protocol):
     def __call__(
         self,
         array: ArrayLike,
-        /,
         indices: _ArrayLikeInt_co,
+        /,
         axis: SupportsIndex,
         dtype: _DTypeLike[_ScalarT],
         out: NDArray[_ScalarT] | None = None,
@@ -1360,7 +1416,7 @@ str_len: _ufunc_1_1[_UFunc11String[np.intp]]
 bitwise_count: Final[_ufunc_1_1] = ...
 
 # {[f]} -> $1
-spacing: Final[_ufunc_1_1] = ...
+spacing: Final[_ufunc_1_1[_Call11Float]] = ...
 
 # {[f]O} -> $1
 cbrt: Final[_ufunc_1_1] = ...

test/static/accept/ufuncs.pyi

@@ -5,9 +5,11 @@ import numpy as np
 import numpy.typing as npt
 
 i8: np.int64
+f4: np.float32
 f8: np.float64
 dt64: np.datetime64
 td64: np.timedelta64
+AR_f4: npt.NDArray[np.float32]
 AR_f8: npt.NDArray[np.float64]
 AR_i8: npt.NDArray[np.int64]
 AR_bool: npt.NDArray[np.bool_]
@@ -88,3 +90,9 @@ assert_type(np.logical_xor(AR_bool, AR_bool, out=AR_bool), npt.NDArray[np.bool_]
 assert_type(np.logical_xor(AR_i8, AR_bool, out=AR_bool), npt.NDArray[np.bool_])
 assert_type(np.logical_xor(AR_bool, AR_i8), npt.NDArray[np.bool_])
 assert_type(np.logical_xor(AR_bool, AR_bool, dtype=np.object_), npt.NDArray[np.object_])
+
+assert_type(np.spacing(f4), np.floating)
+assert_type(np.spacing(f8), np.float64)
+assert_type(np.spacing(AR_f8), npt.NDArray[np.float64])
+assert_type(np.spacing(AR_f4), npt.NDArray[np.floating])
+assert_type(np.spacing(AR_f8, out=AR_f8), npt.NDArray[np.float64])

test/static/reject/ufuncs.pyi

@@ -60,3 +60,6 @@ np.logical_or(dt64, dt64, dtype=np.datetime64)  # type: ignore[arg-type]  # pyri
 
 np.logical_xor(dt64, dt64)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType, reportCallIssue]
 np.logical_xor(dt64, dt64, dtype=np.datetime64)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType, reportCallIssue]
+
+np.spacing(dt64)  # type: ignore[call-overload]  # pyright: ignore[reportCallIssue,reportArgumentType]
+np.spacing(AR_f8, dtype=np.datetime64)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType, reportCallIssue]