1  from typing import (
 2      Literal as L,
 3      Any,
 4      overload,
 5      TypeVar,
 6      Protocol,
 7  )
 8  
 9  from numpy import generic
10  
11  from numpy._typing import (
12      ArrayLike,
13      NDArray,
14      _ArrayLikeInt,
15      _ArrayLike,
16  )
17  
18  _SCT = TypeVar("_SCT", bound=generic)
19  
20  class _ModeFunc(Protocol):
21      def __call__(
22          self,
23          vector: NDArray[Any],
24          iaxis_pad_width: tuple[int, int],
25          iaxis: int,
26          kwargs: dict[str, Any],
27          /,
28      ) -> None: ...
29  
30  _ModeKind = L[
31      "constant",
32      "edge",
33      "linear_ramp",
34      "maximum",
35      "mean",
36      "median",
37      "minimum",
38      "reflect",
39      "symmetric",
40      "wrap",
41      "empty",
42  ]
43  
44  __all__: list[str]
45  
46  # TODO: In practice each keyword argument is exclusive to one or more
47  # specific modes. Consider adding more overloads to express this in the future.
48  
49  # Expand `**kwargs` into explicit keyword-only arguments
50  @overload
51  def pad(
52      array: _ArrayLike[_SCT],
53      pad_width: _ArrayLikeInt,
54      mode: _ModeKind = ...,
55      *,
56      stat_length: None | _ArrayLikeInt = ...,
57      constant_values: ArrayLike = ...,
58      end_values: ArrayLike = ...,
59      reflect_type: L["odd", "even"] = ...,
60  ) -> NDArray[_SCT]: ...
61  @overload
62  def pad(
63      array: ArrayLike,
64      pad_width: _ArrayLikeInt,
65      mode: _ModeKind = ...,
66      *,
67      stat_length: None | _ArrayLikeInt = ...,
68      constant_values: ArrayLike = ...,
69      end_values: ArrayLike = ...,
70      reflect_type: L["odd", "even"] = ...,
71  ) -> NDArray[Any]: ...
72  @overload
73  def pad(
74      array: _ArrayLike[_SCT],
75      pad_width: _ArrayLikeInt,
76      mode: _ModeFunc,
77      **kwargs: Any,
78  ) -> NDArray[_SCT]: ...
79  @overload
80  def pad(
81      array: ArrayLike,
82      pad_width: _ArrayLikeInt,
83      mode: _ModeFunc,
84      **kwargs: Any,
85  ) -> NDArray[Any]: ...