Merge pull request #20036 from BvB93/function_base3

ENH: Add annotations for `np.lib.function_base` part 3/3

2 files changed, 58 insertions, 0 deletions

diff --git a/numpy/typing/tests/data/fail/lib_function_base.pyi b/numpy/typing/tests/data/fail/lib_function_base.pyi
index 2a367b9fe..9cad2da03 100644
--- a/numpy/typing/tests/data/fail/lib_function_base.pyi
+++ b/numpy/typing/tests/data/fail/lib_function_base.pyi
@@ -9,6 +9,8 @@ AR_m: npt.NDArray[np.timedelta64]
 AR_M: npt.NDArray[np.datetime64]
 AR_O: npt.NDArray[np.object_]
 
+def func(a: int) -> None: ...
+
 np.average(AR_m)  # E: incompatible type
 np.select(1, [AR_f8])  # E: incompatible type
 np.angle(AR_m)  # E: incompatible type
@@ -39,3 +41,13 @@ np.hamming(AR_c16)  # E: incompatible type
 np.kaiser(1j, 1)  # E: incompatible type
 np.sinc(AR_O)  # E: incompatible type
 np.median(AR_M)  # E: incompatible type
+
+np.add_newdoc_ufunc(func, "docstring")  # E: incompatible type
+np.percentile(AR_f8, 50j)  # E: No overload variant
+np.percentile(AR_f8, 50, interpolation="bob")  # E: No overload variant
+np.quantile(AR_f8, 0.5j)  # E: No overload variant
+np.quantile(AR_f8, 0.5, interpolation="bob")  # E: No overload variant
+np.meshgrid(AR_f8, AR_f8, indexing="bob")  # E: incompatible type
+np.delete(AR_f8, AR_f8)  # E: incompatible type
+np.insert(AR_f8, AR_f8, 1.5)  # E: incompatible type
+np.digitize(AR_f8, 1j)  # E: No overload variant
diff --git a/numpy/typing/tests/data/reveal/lib_function_base.pyi b/numpy/typing/tests/data/reveal/lib_function_base.pyi
index 81e6d52b2..bced08894 100644
--- a/numpy/typing/tests/data/reveal/lib_function_base.pyi
+++ b/numpy/typing/tests/data/reveal/lib_function_base.pyi
@@ -12,6 +12,7 @@ AR_i8: npt.NDArray[np.int64]
 AR_f8: npt.NDArray[np.float64]
 AR_c16: npt.NDArray[np.complex128]
 AR_m: npt.NDArray[np.timedelta64]
+AR_M: npt.NDArray[np.datetime64]
 AR_O: npt.NDArray[np.object_]
 AR_b: npt.NDArray[np.bool_]
 AR_U: npt.NDArray[np.str_]
@@ -132,3 +133,48 @@ reveal_type(np.median(AR_O))  # E: Any
 reveal_type(np.median(AR_f8, keepdims=True))  # E: Any
 reveal_type(np.median(AR_c16, axis=0))  # E: Any
 reveal_type(np.median(AR_LIKE_f8, out=AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
+
+reveal_type(np.add_newdoc_ufunc(np.add, "docstring"))  # E: None
+
+reveal_type(np.percentile(AR_f8, 50))  # E: numpy.floating[Any]
+reveal_type(np.percentile(AR_c16, 50))  # E: numpy.complexfloating[Any, Any]
+reveal_type(np.percentile(AR_m, 50))  # E: numpy.timedelta64
+reveal_type(np.percentile(AR_M, 50, overwrite_input=True))  # E: numpy.datetime64
+reveal_type(np.percentile(AR_O, 50))  # E: Any
+reveal_type(np.percentile(AR_f8, [50]))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+reveal_type(np.percentile(AR_c16, [50]))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+reveal_type(np.percentile(AR_m, [50]))  # E: numpy.ndarray[Any, numpy.dtype[numpy.timedelta64]]
+reveal_type(np.percentile(AR_M, [50], interpolation="nearest"))  # E: numpy.ndarray[Any, numpy.dtype[numpy.datetime64]]
+reveal_type(np.percentile(AR_O, [50]))  # E: numpy.ndarray[Any, numpy.dtype[numpy.object_]]
+reveal_type(np.percentile(AR_f8, [50], keepdims=True))  # E: Any
+reveal_type(np.percentile(AR_f8, [50], axis=[1]))  # E: Any
+reveal_type(np.percentile(AR_f8, [50], out=AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
+
+reveal_type(np.quantile(AR_f8, 0.5))  # E: numpy.floating[Any]
+reveal_type(np.quantile(AR_c16, 0.5))  # E: numpy.complexfloating[Any, Any]
+reveal_type(np.quantile(AR_m, 0.5))  # E: numpy.timedelta64
+reveal_type(np.quantile(AR_M, 0.5, overwrite_input=True))  # E: numpy.datetime64
+reveal_type(np.quantile(AR_O, 0.5))  # E: Any
+reveal_type(np.quantile(AR_f8, [0.5]))  # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+reveal_type(np.quantile(AR_c16, [0.5]))  # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+reveal_type(np.quantile(AR_m, [0.5]))  # E: numpy.ndarray[Any, numpy.dtype[numpy.timedelta64]]
+reveal_type(np.quantile(AR_M, [0.5], interpolation="nearest"))  # E: numpy.ndarray[Any, numpy.dtype[numpy.datetime64]]
+reveal_type(np.quantile(AR_O, [0.5]))  # E: numpy.ndarray[Any, numpy.dtype[numpy.object_]]
+reveal_type(np.quantile(AR_f8, [0.5], keepdims=True))  # E: Any
+reveal_type(np.quantile(AR_f8, [0.5], axis=[1]))  # E: Any
+reveal_type(np.quantile(AR_f8, [0.5], out=AR_c16))  # E: numpy.ndarray[Any, numpy.dtype[{complex128}]]
+
+reveal_type(np.meshgrid(AR_f8, AR_i8, copy=False))  # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]
+reveal_type(np.meshgrid(AR_f8, AR_i8, AR_c16, indexing="ij"))  # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]
+
+reveal_type(np.delete(AR_f8, np.s_[:5]))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.delete(AR_LIKE_f8, [0, 4, 9], axis=0))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.insert(AR_f8, np.s_[:5], 5))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.insert(AR_LIKE_f8, [0, 4, 9], [0.5, 9.2, 7], axis=0))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.append(AR_f8, 5))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.append(AR_LIKE_f8, 1j, axis=0))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.digitize(4.5, [1]))  # E: {intp}
+reveal_type(np.digitize(AR_f8, [1, 2, 3]))  # E: numpy.ndarray[Any, numpy.dtype[{intp}]]