shape_base.pyi 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220
  1. import sys
  2. from collections.abc import Callable, Sequence
  3. from typing import TypeVar, Any, overload, SupportsIndex, Protocol
  4. if sys.version_info >= (3, 10):
  5. from typing import ParamSpec, Concatenate
  6. else:
  7. from typing_extensions import ParamSpec, Concatenate
  8. from numpy import (
  9. generic,
  10. integer,
  11. ufunc,
  12. bool_,
  13. unsignedinteger,
  14. signedinteger,
  15. floating,
  16. complexfloating,
  17. object_,
  18. )
  19. from numpy._typing import (
  20. ArrayLike,
  21. NDArray,
  22. _ShapeLike,
  23. _ArrayLike,
  24. _ArrayLikeBool_co,
  25. _ArrayLikeUInt_co,
  26. _ArrayLikeInt_co,
  27. _ArrayLikeFloat_co,
  28. _ArrayLikeComplex_co,
  29. _ArrayLikeObject_co,
  30. )
  31. from numpy.core.shape_base import vstack
  32. _P = ParamSpec("_P")
  33. _SCT = TypeVar("_SCT", bound=generic)
  34. # The signatures of `__array_wrap__` and `__array_prepare__` are the same;
  35. # give them unique names for the sake of clarity
  36. class _ArrayWrap(Protocol):
  37. def __call__(
  38. self,
  39. array: NDArray[Any],
  40. context: None | tuple[ufunc, tuple[Any, ...], int] = ...,
  41. /,
  42. ) -> Any: ...
  43. class _ArrayPrepare(Protocol):
  44. def __call__(
  45. self,
  46. array: NDArray[Any],
  47. context: None | tuple[ufunc, tuple[Any, ...], int] = ...,
  48. /,
  49. ) -> Any: ...
  50. class _SupportsArrayWrap(Protocol):
  51. @property
  52. def __array_wrap__(self) -> _ArrayWrap: ...
  53. class _SupportsArrayPrepare(Protocol):
  54. @property
  55. def __array_prepare__(self) -> _ArrayPrepare: ...
  56. __all__: list[str]
  57. row_stack = vstack
  58. def take_along_axis(
  59. arr: _SCT | NDArray[_SCT],
  60. indices: NDArray[integer[Any]],
  61. axis: None | int,
  62. ) -> NDArray[_SCT]: ...
  63. def put_along_axis(
  64. arr: NDArray[_SCT],
  65. indices: NDArray[integer[Any]],
  66. values: ArrayLike,
  67. axis: None | int,
  68. ) -> None: ...
  69. @overload
  70. def apply_along_axis(
  71. func1d: Callable[Concatenate[NDArray[Any], _P], _ArrayLike[_SCT]],
  72. axis: SupportsIndex,
  73. arr: ArrayLike,
  74. *args: _P.args,
  75. **kwargs: _P.kwargs,
  76. ) -> NDArray[_SCT]: ...
  77. @overload
  78. def apply_along_axis(
  79. func1d: Callable[Concatenate[NDArray[Any], _P], ArrayLike],
  80. axis: SupportsIndex,
  81. arr: ArrayLike,
  82. *args: _P.args,
  83. **kwargs: _P.kwargs,
  84. ) -> NDArray[Any]: ...
  85. def apply_over_axes(
  86. func: Callable[[NDArray[Any], int], NDArray[_SCT]],
  87. a: ArrayLike,
  88. axes: int | Sequence[int],
  89. ) -> NDArray[_SCT]: ...
  90. @overload
  91. def expand_dims(
  92. a: _ArrayLike[_SCT],
  93. axis: _ShapeLike,
  94. ) -> NDArray[_SCT]: ...
  95. @overload
  96. def expand_dims(
  97. a: ArrayLike,
  98. axis: _ShapeLike,
  99. ) -> NDArray[Any]: ...
  100. @overload
  101. def column_stack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ...
  102. @overload
  103. def column_stack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
  104. @overload
  105. def dstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ...
  106. @overload
  107. def dstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ...
  108. @overload
  109. def array_split(
  110. ary: _ArrayLike[_SCT],
  111. indices_or_sections: _ShapeLike,
  112. axis: SupportsIndex = ...,
  113. ) -> list[NDArray[_SCT]]: ...
  114. @overload
  115. def array_split(
  116. ary: ArrayLike,
  117. indices_or_sections: _ShapeLike,
  118. axis: SupportsIndex = ...,
  119. ) -> list[NDArray[Any]]: ...
  120. @overload
  121. def split(
  122. ary: _ArrayLike[_SCT],
  123. indices_or_sections: _ShapeLike,
  124. axis: SupportsIndex = ...,
  125. ) -> list[NDArray[_SCT]]: ...
  126. @overload
  127. def split(
  128. ary: ArrayLike,
  129. indices_or_sections: _ShapeLike,
  130. axis: SupportsIndex = ...,
  131. ) -> list[NDArray[Any]]: ...
  132. @overload
  133. def hsplit(
  134. ary: _ArrayLike[_SCT],
  135. indices_or_sections: _ShapeLike,
  136. ) -> list[NDArray[_SCT]]: ...
  137. @overload
  138. def hsplit(
  139. ary: ArrayLike,
  140. indices_or_sections: _ShapeLike,
  141. ) -> list[NDArray[Any]]: ...
  142. @overload
  143. def vsplit(
  144. ary: _ArrayLike[_SCT],
  145. indices_or_sections: _ShapeLike,
  146. ) -> list[NDArray[_SCT]]: ...
  147. @overload
  148. def vsplit(
  149. ary: ArrayLike,
  150. indices_or_sections: _ShapeLike,
  151. ) -> list[NDArray[Any]]: ...
  152. @overload
  153. def dsplit(
  154. ary: _ArrayLike[_SCT],
  155. indices_or_sections: _ShapeLike,
  156. ) -> list[NDArray[_SCT]]: ...
  157. @overload
  158. def dsplit(
  159. ary: ArrayLike,
  160. indices_or_sections: _ShapeLike,
  161. ) -> list[NDArray[Any]]: ...
  162. @overload
  163. def get_array_prepare(*args: _SupportsArrayPrepare) -> _ArrayPrepare: ...
  164. @overload
  165. def get_array_prepare(*args: object) -> None | _ArrayPrepare: ...
  166. @overload
  167. def get_array_wrap(*args: _SupportsArrayWrap) -> _ArrayWrap: ...
  168. @overload
  169. def get_array_wrap(*args: object) -> None | _ArrayWrap: ...
  170. @overload
  171. def kron(a: _ArrayLikeBool_co, b: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc]
  172. @overload
  173. def kron(a: _ArrayLikeUInt_co, b: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc]
  174. @overload
  175. def kron(a: _ArrayLikeInt_co, b: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc]
  176. @overload
  177. def kron(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc]
  178. @overload
  179. def kron(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ...
  180. @overload
  181. def kron(a: _ArrayLikeObject_co, b: Any) -> NDArray[object_]: ...
  182. @overload
  183. def kron(a: Any, b: _ArrayLikeObject_co) -> NDArray[object_]: ...
  184. @overload
  185. def tile(
  186. A: _ArrayLike[_SCT],
  187. reps: int | Sequence[int],
  188. ) -> NDArray[_SCT]: ...
  189. @overload
  190. def tile(
  191. A: ArrayLike,
  192. reps: int | Sequence[int],
  193. ) -> NDArray[Any]: ...