TYP: Fix interp to accept and return scalars

Co-authored-by: PTUsumit <[email protected]>

Author: jorenham

numpy/lib/_function_base_impl.pyi

@@ -1,4 +1,4 @@
-from collections.abc import Sequence, Iterator, Callable, Iterable
+from collections.abc import Sequence, Callable, Iterable
 from typing import (
     Concatenate,
     Literal as L,
@@ -15,8 +15,9 @@ from typing import (
 )
 from typing_extensions import deprecated
 
+import numpy as np
 from numpy import (
-    vectorize as vectorize,
+    vectorize,
     generic,
     integer,
     floating,
@@ -35,19 +36,22 @@ from numpy._typing import (
     NDArray,
     ArrayLike,
     DTypeLike,
-    _ShapeLike,
-    _ScalarLike_co,
-    _DTypeLike,
     _ArrayLike,
+    _DTypeLike,
+    _ShapeLike,
     _ArrayLikeBool_co,
     _ArrayLikeInt_co,
     _ArrayLikeFloat_co,
     _ArrayLikeComplex_co,
+    _ArrayLikeNumber_co,
     _ArrayLikeTD64_co,
     _ArrayLikeDT64_co,
     _ArrayLikeObject_co,
     _FloatLike_co,
     _ComplexLike_co,
+    _NumberLike_co,
+    _ScalarLike_co,
+    _NestedSequence
 )
 
 __all__ = [
@@ -303,24 +307,87 @@ def diff(
     append: ArrayLike = ...,
 ) -> NDArray[Any]: ...
 
-@overload
+@overload  # float scalar
 def interp(
-    x: _ArrayLikeFloat_co,
+    x: _FloatLike_co,
+    xp: _ArrayLikeFloat_co,
+    fp: _ArrayLikeFloat_co,
+    left: _FloatLike_co | None = None,
+    right: _FloatLike_co | None = None,
+    period: _FloatLike_co | None = None,
+) -> float64: ...
+@overload  # float array
+def interp(
+    x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co],
     xp: _ArrayLikeFloat_co,
     fp: _ArrayLikeFloat_co,
-    left: None | _FloatLike_co = ...,
-    right: None | _FloatLike_co = ...,
-    period: None | _FloatLike_co = ...,
+    left: _FloatLike_co | None = None,
+    right: _FloatLike_co | None = None,
+    period: _FloatLike_co | None = None,
 ) -> NDArray[float64]: ...
-@overload
+@overload  # float scalar or array
 def interp(
     x: _ArrayLikeFloat_co,
     xp: _ArrayLikeFloat_co,
-    fp: _ArrayLikeComplex_co,
-    left: None | _ComplexLike_co = ...,
-    right: None | _ComplexLike_co = ...,
-    period: None | _FloatLike_co = ...,
+    fp: _ArrayLikeFloat_co,
+    left: _FloatLike_co | None = None,
+    right: _FloatLike_co | None = None,
+    period: _FloatLike_co | None = None,
+) -> NDArray[float64] | float64: ...
+@overload  # complex scalar
+def interp(
+    x: _FloatLike_co,
+    xp: _ArrayLikeFloat_co,
+    fp: _ArrayLike[complexfloating],
+    left: _NumberLike_co | None = None,
+    right: _NumberLike_co | None = None,
+    period: _FloatLike_co | None = None,
+) -> complex128: ...
+@overload  # complex or float scalar
+def interp(
+    x: _FloatLike_co,
+    xp: _ArrayLikeFloat_co,
+    fp: Sequence[complex | complexfloating],
+    left: _NumberLike_co | None = None,
+    right: _NumberLike_co | None = None,
+    period: _FloatLike_co | None = None,
+) -> complex128 | float64: ...
+@overload  # complex array
+def interp(
+    x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co],
+    xp: _ArrayLikeFloat_co,
+    fp: _ArrayLike[complexfloating],
+    left: _NumberLike_co | None = None,
+    right: _NumberLike_co | None = None,
+    period: _FloatLike_co | None = None,
 ) -> NDArray[complex128]: ...
+@overload  # complex or float array
+def interp(
+    x: NDArray[floating | integer | np.bool] | _NestedSequence[_FloatLike_co],
+    xp: _ArrayLikeFloat_co,
+    fp: Sequence[complex | complexfloating],
+    left: _NumberLike_co | None = None,
+    right: _NumberLike_co | None = None,
+    period: _FloatLike_co | None = None,
+) -> NDArray[complex128 | float64]: ...
+@overload  # complex scalar or array
+def interp(
+    x: _ArrayLikeFloat_co,
+    xp: _ArrayLikeFloat_co,
+    fp: _ArrayLike[complexfloating],
+    left: _NumberLike_co | None = None,
+    right: _NumberLike_co | None = None,
+    period: _FloatLike_co | None = None,
+) -> NDArray[complex128] | complex128: ...
+@overload  # complex or float scalar or array
+def interp(
+    x: _ArrayLikeFloat_co,
+    xp: _ArrayLikeFloat_co,
+    fp: _ArrayLikeNumber_co,
+    left: _NumberLike_co | None = None,
+    right: _NumberLike_co | None = None,
+    period: _FloatLike_co | None = None,
+) -> NDArray[complex128 | float64] | complex128 | float64: ...
 
 @overload
 def angle(z: _ComplexLike_co, deg: bool = ...) -> floating[Any]: ...

numpy/typing/tests/data/reveal/lib_function_base.pyi

@@ -94,6 +94,15 @@ assert_type(np.diff("bob", n=0), str)
 assert_type(np.diff(AR_f8, axis=0), npt.NDArray[Any])
 assert_type(np.diff(AR_LIKE_f8, prepend=1.5), npt.NDArray[Any])
 
+assert_type(np.interp(1, [1], AR_f8), np.float64)
+assert_type(np.interp(1, [1], [1]), np.float64)
+assert_type(np.interp(1, [1], AR_c16), np.complex128)
+assert_type(np.interp(1, [1], [1j]), np.complex128)  # pyright correctly infers `complex128 | float64`
+assert_type(np.interp([1], [1], AR_f8), npt.NDArray[np.float64])
+assert_type(np.interp([1], [1], [1]), npt.NDArray[np.float64])
+assert_type(np.interp([1], [1], AR_c16), npt.NDArray[np.complex128])
+assert_type(np.interp([1], [1], [1j]), npt.NDArray[np.complex128])  # pyright correctly infers `NDArray[complex128 | float64]`
+
 assert_type(np.angle(f8), np.floating[Any])
 assert_type(np.angle(AR_f8), npt.NDArray[np.floating[Any]])
 assert_type(np.angle(AR_c16, deg=True), npt.NDArray[np.floating[Any]])