shape_base.pyi 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. from collections.abc import Sequence
  2. from typing import Any, SupportsIndex, TypeVar, overload
  3. from numpy import _CastingKind, generic
  4. from numpy._typing import ArrayLike, DTypeLike, NDArray, _ArrayLike, _DTypeLike
  5. __all__ = [
  6. "atleast_1d",
  7. "atleast_2d",
  8. "atleast_3d",
  9. "block",
  10. "hstack",
  11. "stack",
  12. "unstack",
  13. "vstack",
  14. ]
  15. _SCT = TypeVar("_SCT", bound=generic)
  16. _SCT1 = TypeVar("_SCT1", bound=generic)
  17. _SCT2 = TypeVar("_SCT2", bound=generic)
  18. _ArrayT = TypeVar("_ArrayT", bound=NDArray[Any])
  19. ###
  20. @overload
  21. def atleast_1d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
  22. @overload
  23. def atleast_1d(a0: _ArrayLike[_SCT1], a1: _ArrayLike[_SCT2], /) -> tuple[NDArray[_SCT1], NDArray[_SCT2]]: ...
  24. @overload
  25. def atleast_1d(a0: _ArrayLike[_SCT], a1: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ...
  26. @overload
  27. def atleast_1d(a0: ArrayLike, /) -> NDArray[Any]: ...
  28. @overload
  29. def atleast_1d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
  30. @overload
  31. def atleast_1d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ...
  32. #
  33. @overload
  34. def atleast_2d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
  35. @overload
  36. def atleast_2d(a0: _ArrayLike[_SCT1], a1: _ArrayLike[_SCT2], /) -> tuple[NDArray[_SCT1], NDArray[_SCT2]]: ...
  37. @overload
  38. def atleast_2d(a0: _ArrayLike[_SCT], a1: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ...
  39. @overload
  40. def atleast_2d(a0: ArrayLike, /) -> NDArray[Any]: ...
  41. @overload
  42. def atleast_2d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
  43. @overload
  44. def atleast_2d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ...
  45. #
  46. @overload
  47. def atleast_3d(a0: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
  48. @overload
  49. def atleast_3d(a0: _ArrayLike[_SCT1], a1: _ArrayLike[_SCT2], /) -> tuple[NDArray[_SCT1], NDArray[_SCT2]]: ...
  50. @overload
  51. def atleast_3d(a0: _ArrayLike[_SCT], a1: _ArrayLike[_SCT], /, *arys: _ArrayLike[_SCT]) -> tuple[NDArray[_SCT], ...]: ...
  52. @overload
  53. def atleast_3d(a0: ArrayLike, /) -> NDArray[Any]: ...
  54. @overload
  55. def atleast_3d(a0: ArrayLike, a1: ArrayLike, /) -> tuple[NDArray[Any], NDArray[Any]]: ...
  56. @overload
  57. def atleast_3d(a0: ArrayLike, a1: ArrayLike, /, *ai: ArrayLike) -> tuple[NDArray[Any], ...]: ...
  58. #
  59. @overload
  60. def vstack(
  61. tup: Sequence[_ArrayLike[_SCT]],
  62. *,
  63. dtype: None = ...,
  64. casting: _CastingKind = ...
  65. ) -> NDArray[_SCT]: ...
  66. @overload
  67. def vstack(
  68. tup: Sequence[ArrayLike],
  69. *,
  70. dtype: _DTypeLike[_SCT],
  71. casting: _CastingKind = ...
  72. ) -> NDArray[_SCT]: ...
  73. @overload
  74. def vstack(
  75. tup: Sequence[ArrayLike],
  76. *,
  77. dtype: DTypeLike = ...,
  78. casting: _CastingKind = ...
  79. ) -> NDArray[Any]: ...
  80. @overload
  81. def hstack(
  82. tup: Sequence[_ArrayLike[_SCT]],
  83. *,
  84. dtype: None = ...,
  85. casting: _CastingKind = ...
  86. ) -> NDArray[_SCT]: ...
  87. @overload
  88. def hstack(
  89. tup: Sequence[ArrayLike],
  90. *,
  91. dtype: _DTypeLike[_SCT],
  92. casting: _CastingKind = ...
  93. ) -> NDArray[_SCT]: ...
  94. @overload
  95. def hstack(
  96. tup: Sequence[ArrayLike],
  97. *,
  98. dtype: DTypeLike = ...,
  99. casting: _CastingKind = ...
  100. ) -> NDArray[Any]: ...
  101. @overload
  102. def stack(
  103. arrays: Sequence[_ArrayLike[_SCT]],
  104. axis: SupportsIndex = ...,
  105. out: None = ...,
  106. *,
  107. dtype: None = ...,
  108. casting: _CastingKind = ...
  109. ) -> NDArray[_SCT]: ...
  110. @overload
  111. def stack(
  112. arrays: Sequence[ArrayLike],
  113. axis: SupportsIndex = ...,
  114. out: None = ...,
  115. *,
  116. dtype: _DTypeLike[_SCT],
  117. casting: _CastingKind = ...
  118. ) -> NDArray[_SCT]: ...
  119. @overload
  120. def stack(
  121. arrays: Sequence[ArrayLike],
  122. axis: SupportsIndex = ...,
  123. out: None = ...,
  124. *,
  125. dtype: DTypeLike = ...,
  126. casting: _CastingKind = ...
  127. ) -> NDArray[Any]: ...
  128. @overload
  129. def stack(
  130. arrays: Sequence[ArrayLike],
  131. axis: SupportsIndex,
  132. out: _ArrayT,
  133. *,
  134. dtype: DTypeLike | None = None,
  135. casting: _CastingKind = "same_kind",
  136. ) -> _ArrayT: ...
  137. @overload
  138. def stack(
  139. arrays: Sequence[ArrayLike],
  140. axis: SupportsIndex = 0,
  141. *,
  142. out: _ArrayT,
  143. dtype: DTypeLike | None = None,
  144. casting: _CastingKind = "same_kind",
  145. ) -> _ArrayT: ...
  146. @overload
  147. def unstack(
  148. array: _ArrayLike[_SCT],
  149. /,
  150. *,
  151. axis: int = ...,
  152. ) -> tuple[NDArray[_SCT], ...]: ...
  153. @overload
  154. def unstack(
  155. array: ArrayLike,
  156. /,
  157. *,
  158. axis: int = ...,
  159. ) -> tuple[NDArray[Any], ...]: ...
  160. @overload
  161. def block(arrays: _ArrayLike[_SCT]) -> NDArray[_SCT]: ...
  162. @overload
  163. def block(arrays: ArrayLike) -> NDArray[Any]: ...