arraypad.pyi 1.7 KB

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