codekingpro/portable-devtools

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