Merge pull request #19223 from BvB93/arraypad

ENH: Add annotations for `np.pad`

3 files changed, 119 insertions, 2 deletions

diff --git a/numpy/lib/arraypad.pyi b/numpy/lib/arraypad.pyi
index 64e3e1331..df9538dd7 100644
--- a/numpy/lib/arraypad.pyi
+++ b/numpy/lib/arraypad.pyi
@@ -1,5 +1,94 @@
-from typing import List
+import sys
+from typing import (
+    Any,
+    Dict,
+    List,
+    overload,
+    Tuple,
+    TypeVar,
+)
+
+from numpy import ndarray, dtype, generic
+
+from numpy.typing import (
+    ArrayLike,
+    NDArray,
+    _ArrayLikeInt,
+    _NestedSequence,
+    _SupportsArray,
+)
+
+if sys.version_info >= (3, 8):
+    from typing import Literal as L, Protocol
+else:
+    from typing_extensions import Literal as L, Protocol
+
+_SCT = TypeVar("_SCT", bound=generic)
+
+class _ModeFunc(Protocol):
+    def __call__(
+        self,
+        __vector: NDArray[Any],
+        __iaxis_pad_width: Tuple[int, int],
+        __iaxis: int,
+        __kwargs: Dict[str, Any],
+    ) -> None: ...
+
+_ModeKind = L[
+    "constant",
+    "edge",
+    "linear_ramp",
+    "maximum",
+    "mean",
+    "median",
+    "minimum",
+    "reflect",
+    "symmetric",
+    "wrap",
+    "empty",
+]
+
+_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
 
 __all__: List[str]
 
-def pad(array, pad_width, mode=..., **kwargs): ...
+# TODO: In practice each keyword argument is exclusive to one or more
+# specific modes. Consider adding more overloads to express this in the future.
+
+# Expand `**kwargs` into explicit keyword-only arguments
+@overload
+def pad(
+    array: _ArrayLike[_SCT],
+    pad_width: _ArrayLikeInt,
+    mode: _ModeKind = ...,
+    *,
+    stat_length: None | _ArrayLikeInt = ...,
+    constant_values: ArrayLike = ...,
+    end_values: ArrayLike = ...,
+    reflect_type: L["odd", "even"] = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def pad(
+    array: ArrayLike,
+    pad_width: _ArrayLikeInt,
+    mode: _ModeKind = ...,
+    *,
+    stat_length: None | _ArrayLikeInt = ...,
+    constant_values: ArrayLike = ...,
+    end_values: ArrayLike = ...,
+    reflect_type: L["odd", "even"] = ...,
+) -> NDArray[Any]: ...
+@overload
+def pad(
+    array: _ArrayLike[_SCT],
+    pad_width: _ArrayLikeInt,
+    mode: _ModeFunc,
+    **kwargs: Any,
+) -> NDArray[_SCT]: ...
+@overload
+def pad(
+    array: ArrayLike,
+    pad_width: _ArrayLikeInt,
+    mode: _ModeFunc,
+    **kwargs: Any,
+) -> NDArray[Any]: ...
diff --git a/numpy/typing/tests/data/fail/array_pad.py b/numpy/typing/tests/data/fail/array_pad.py
new file mode 100644
index 000000000..2be51a871
--- /dev/null
+++ b/numpy/typing/tests/data/fail/array_pad.py
@@ -0,0 +1,6 @@
+import numpy as np
+import numpy.typing as npt
+
+AR_i8: npt.NDArray[np.int64]
+
+np.pad(AR_i8, 2, mode="bob")  # E: No overload variant
diff --git a/numpy/typing/tests/data/reveal/arraypad.py b/numpy/typing/tests/data/reveal/arraypad.py
new file mode 100644
index 000000000..ba5577ee0
--- /dev/null
+++ b/numpy/typing/tests/data/reveal/arraypad.py
@@ -0,0 +1,22 @@
+from typing import List, Any, Mapping, Tuple
+from typing_extensions import SupportsIndex
+
+import numpy as np
+import numpy.typing as npt
+
+def mode_func(
+    ar: npt.NDArray[np.number[Any]],
+    width: Tuple[int, int],
+    iaxis: SupportsIndex,
+    kwargs: Mapping[str, Any],
+) -> None: ...
+
+AR_i8: npt.NDArray[np.int64]
+AR_f8: npt.NDArray[np.float64]
+AR_LIKE: List[int]
+
+reveal_type(np.pad(AR_i8, (2, 3), "constant"))  # E: numpy.ndarray[Any, numpy.dtype[{int64}]]
+reveal_type(np.pad(AR_LIKE, (2, 3), "constant"))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+
+reveal_type(np.pad(AR_f8, (2, 3), mode_func))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.pad(AR_f8, (2, 3), mode_func, a=1, b=2))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]