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

1 files changed, 91 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]: ...