Merge pull request #18377 from BvB93/ufunclike

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

3 files changed, 60 insertions, 5 deletions

diff --git a/numpy/lib/__init__.pyi b/numpy/lib/__init__.pyi
index a8eb24207..4468d27e9 100644
--- a/numpy/lib/__init__.pyi
+++ b/numpy/lib/__init__.pyi
@@ -1,5 +1,11 @@
 from typing import Any, List
 
+from numpy.lib.ufunclike import (
+    fix as fix,
+    isposinf as isposinf,
+    isneginf as isneginf,
+)
+
 __all__: List[str]
 
 emath: Any
@@ -108,9 +114,6 @@ tril_indices: Any
 tril_indices_from: Any
 triu_indices: Any
 triu_indices_from: Any
-fix: Any
-isneginf: Any
-isposinf: Any
 pad: Any
 poly: Any
 roots: Any
diff --git a/numpy/lib/ufunclike.py b/numpy/lib/ufunclike.py
index 1f26a1845..0956de82b 100644
--- a/numpy/lib/ufunclike.py
+++ b/numpy/lib/ufunclike.py
@@ -189,7 +189,8 @@ def isposinf(x, out=None):
     try:
         signbit = ~nx.signbit(x)
     except TypeError as e:
-        raise TypeError('This operation is not supported for complex values '
+        dtype = nx.asanyarray(x).dtype
+        raise TypeError(f'This operation is not supported for {dtype} values '
                         'because it would be ambiguous.') from e
     else:
         return nx.logical_and(is_inf, signbit, out)
@@ -260,7 +261,8 @@ def isneginf(x, out=None):
     try:
         signbit = nx.signbit(x)
     except TypeError as e:
-        raise TypeError('This operation is not supported for complex values '
+        dtype = nx.asanyarray(x).dtype
+        raise TypeError(f'This operation is not supported for {dtype} values '
                         'because it would be ambiguous.') from e
     else:
         return nx.logical_and(is_inf, signbit, out)
diff --git a/numpy/lib/ufunclike.pyi b/numpy/lib/ufunclike.pyi
new file mode 100644
index 000000000..3443fa7ae
--- /dev/null
+++ b/numpy/lib/ufunclike.pyi
@@ -0,0 +1,50 @@
+from typing import Any, overload, TypeVar, List, Union
+
+from numpy import floating, bool_, ndarray
+from numpy.typing import (
+    _ArrayLikeFloat_co,
+    _ArrayLikeObject_co,
+    _ArrayOrScalar,
+)
+
+_ArrayType = TypeVar("_ArrayType", bound=ndarray[Any, Any])
+
+__all__: List[str]
+
+@overload
+def fix(
+    x: _ArrayLikeFloat_co,
+    out: None = ...,
+) -> _ArrayOrScalar[floating[Any]]: ...
+@overload
+def fix(
+    x: _ArrayLikeObject_co,
+    out: None = ...,
+) -> Any: ...
+@overload
+def fix(
+    x: Union[_ArrayLikeFloat_co, _ArrayLikeObject_co],
+    out: _ArrayType,
+) -> _ArrayType: ...
+
+@overload
+def isposinf(
+    x: _ArrayLikeFloat_co,
+    out: None = ...,
+) -> _ArrayOrScalar[bool_]: ...
+@overload
+def isposinf(
+    x: _ArrayLikeFloat_co,
+    out: _ArrayType,
+) -> _ArrayType: ...
+
+@overload
+def isneginf(
+    x: _ArrayLikeFloat_co,
+    out: None = ...,
+) -> _ArrayOrScalar[bool_]: ...
+@overload
+def isneginf(
+    x: _ArrayLikeFloat_co,
+    out: _ArrayType,
+) -> _ArrayType: ...