summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--numpy/lib/arraypad.pyi93
-rw-r--r--numpy/typing/tests/data/fail/array_pad.py6
-rw-r--r--numpy/typing/tests/data/reveal/arraypad.py22
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}]]