_arraypad_impl.pyi 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. from typing import (
  2. Literal as L,
  3. Any,
  4. TypeAlias,
  5. overload,
  6. TypeVar,
  7. Protocol,
  8. type_check_only,
  9. )
  10. from numpy import generic
  11. from numpy._typing import (
  12. ArrayLike,
  13. NDArray,
  14. _ArrayLikeInt,
  15. _ArrayLike,
  16. )
  17. __all__ = ["pad"]
  18. _SCT = TypeVar("_SCT", bound=generic)
  19. @type_check_only
  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. _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. # TODO: In practice each keyword argument is exclusive to one or more
  43. # specific modes. Consider adding more overloads to express this in the future.
  44. # Expand `**kwargs` into explicit keyword-only arguments
  45. @overload
  46. def pad(
  47. array: _ArrayLike[_SCT],
  48. pad_width: _ArrayLikeInt,
  49. mode: _ModeKind = ...,
  50. *,
  51. stat_length: None | _ArrayLikeInt = ...,
  52. constant_values: ArrayLike = ...,
  53. end_values: ArrayLike = ...,
  54. reflect_type: L["odd", "even"] = ...,
  55. ) -> NDArray[_SCT]: ...
  56. @overload
  57. def pad(
  58. array: ArrayLike,
  59. pad_width: _ArrayLikeInt,
  60. mode: _ModeKind = ...,
  61. *,
  62. stat_length: None | _ArrayLikeInt = ...,
  63. constant_values: ArrayLike = ...,
  64. end_values: ArrayLike = ...,
  65. reflect_type: L["odd", "even"] = ...,
  66. ) -> NDArray[Any]: ...
  67. @overload
  68. def pad(
  69. array: _ArrayLike[_SCT],
  70. pad_width: _ArrayLikeInt,
  71. mode: _ModeFunc,
  72. **kwargs: Any,
  73. ) -> NDArray[_SCT]: ...
  74. @overload
  75. def pad(
  76. array: ArrayLike,
  77. pad_width: _ArrayLikeInt,
  78. mode: _ModeFunc,
  79. **kwargs: Any,
  80. ) -> NDArray[Any]: ...