fromnumeric.pyi 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357
  1. import sys
  2. import datetime as dt
  3. from typing import Optional, Union, Sequence, Tuple, Any, overload, TypeVar
  4. from numpy import (
  5. ndarray,
  6. signedinteger,
  7. bool_,
  8. generic,
  9. _OrderKACF,
  10. _OrderACF,
  11. _ArrayLikeBool,
  12. _ArrayLikeIntOrBool,
  13. _ModeKind,
  14. _PartitionKind,
  15. _SortKind,
  16. _SortSide,
  17. )
  18. from numpy.typing import (
  19. DTypeLike,
  20. ArrayLike,
  21. _ShapeLike,
  22. _Shape,
  23. _NumberLike,
  24. )
  25. if sys.version_info >= (3, 8):
  26. from typing import Literal
  27. else:
  28. from typing_extensions import Literal
  29. # Various annotations for scalars
  30. # While dt.datetime and dt.timedelta are not technically part of NumPy,
  31. # they are one of the rare few builtin scalars which serve as valid return types.
  32. # See https://github.com/numpy/numpy-stubs/pull/67#discussion_r412604113.
  33. _ScalarNumpy = Union[generic, dt.datetime, dt.timedelta]
  34. _ScalarBuiltin = Union[str, bytes, dt.date, dt.timedelta, bool, int, float, complex]
  35. _Scalar = Union[_ScalarBuiltin, _ScalarNumpy]
  36. # Integers and booleans can generally be used interchangeably
  37. _ScalarGeneric = TypeVar("_ScalarGeneric", bound=generic)
  38. # The signature of take() follows a common theme with its overloads:
  39. # 1. A generic comes in; the same generic comes out
  40. # 2. A scalar comes in; a generic comes out
  41. # 3. An array-like object comes in; some keyword ensures that a generic comes out
  42. # 4. An array-like object comes in; an ndarray or generic comes out
  43. def take(
  44. a: ArrayLike,
  45. indices: _ArrayLikeIntOrBool,
  46. axis: Optional[int] = ...,
  47. out: Optional[ndarray] = ...,
  48. mode: _ModeKind = ...,
  49. ) -> Any: ...
  50. def reshape(
  51. a: ArrayLike,
  52. newshape: _ShapeLike,
  53. order: _OrderACF = ...,
  54. ) -> ndarray: ...
  55. def choose(
  56. a: _ArrayLikeIntOrBool,
  57. choices: ArrayLike,
  58. out: Optional[ndarray] = ...,
  59. mode: _ModeKind = ...,
  60. ) -> Any: ...
  61. def repeat(
  62. a: ArrayLike,
  63. repeats: _ArrayLikeIntOrBool,
  64. axis: Optional[int] = ...,
  65. ) -> ndarray: ...
  66. def put(
  67. a: ndarray,
  68. ind: _ArrayLikeIntOrBool,
  69. v: ArrayLike,
  70. mode: _ModeKind = ...,
  71. ) -> None: ...
  72. def swapaxes(
  73. a: ArrayLike,
  74. axis1: int,
  75. axis2: int,
  76. ) -> ndarray: ...
  77. def transpose(
  78. a: ArrayLike,
  79. axes: Union[None, Sequence[int], ndarray] = ...
  80. ) -> ndarray: ...
  81. def partition(
  82. a: ArrayLike,
  83. kth: _ArrayLikeIntOrBool,
  84. axis: Optional[int] = ...,
  85. kind: _PartitionKind = ...,
  86. order: Union[None, str, Sequence[str]] = ...,
  87. ) -> ndarray: ...
  88. def argpartition(
  89. a: ArrayLike,
  90. kth: _ArrayLikeIntOrBool,
  91. axis: Optional[int] = ...,
  92. kind: _PartitionKind = ...,
  93. order: Union[None, str, Sequence[str]] = ...,
  94. ) -> Any: ...
  95. def sort(
  96. a: ArrayLike,
  97. axis: Optional[int] = ...,
  98. kind: Optional[_SortKind] = ...,
  99. order: Union[None, str, Sequence[str]] = ...,
  100. ) -> ndarray: ...
  101. def argsort(
  102. a: ArrayLike,
  103. axis: Optional[int] = ...,
  104. kind: Optional[_SortKind] = ...,
  105. order: Union[None, str, Sequence[str]] = ...,
  106. ) -> ndarray: ...
  107. @overload
  108. def argmax(
  109. a: ArrayLike,
  110. axis: None = ...,
  111. out: Optional[ndarray] = ...,
  112. ) -> signedinteger[Any]: ...
  113. @overload
  114. def argmax(
  115. a: ArrayLike,
  116. axis: Optional[int] = ...,
  117. out: Optional[ndarray] = ...,
  118. ) -> Any: ...
  119. @overload
  120. def argmin(
  121. a: ArrayLike,
  122. axis: None = ...,
  123. out: Optional[ndarray] = ...,
  124. ) -> signedinteger[Any]: ...
  125. @overload
  126. def argmin(
  127. a: ArrayLike,
  128. axis: Optional[int] = ...,
  129. out: Optional[ndarray] = ...,
  130. ) -> Any: ...
  131. @overload
  132. def searchsorted(
  133. a: ArrayLike,
  134. v: _Scalar,
  135. side: _SortSide = ...,
  136. sorter: Optional[_ArrayLikeIntOrBool] = ..., # 1D int array
  137. ) -> signedinteger[Any]: ...
  138. @overload
  139. def searchsorted(
  140. a: ArrayLike,
  141. v: ArrayLike,
  142. side: _SortSide = ...,
  143. sorter: Optional[_ArrayLikeIntOrBool] = ..., # 1D int array
  144. ) -> ndarray: ...
  145. def resize(
  146. a: ArrayLike,
  147. new_shape: _ShapeLike,
  148. ) -> ndarray: ...
  149. @overload
  150. def squeeze(
  151. a: _ScalarGeneric,
  152. axis: Optional[_ShapeLike] = ...,
  153. ) -> _ScalarGeneric: ...
  154. @overload
  155. def squeeze(
  156. a: ArrayLike,
  157. axis: Optional[_ShapeLike] = ...,
  158. ) -> ndarray: ...
  159. def diagonal(
  160. a: ArrayLike,
  161. offset: int = ...,
  162. axis1: int = ...,
  163. axis2: int = ..., # >= 2D array
  164. ) -> ndarray: ...
  165. def trace(
  166. a: ArrayLike, # >= 2D array
  167. offset: int = ...,
  168. axis1: int = ...,
  169. axis2: int = ...,
  170. dtype: DTypeLike = ...,
  171. out: Optional[ndarray] = ...,
  172. ) -> Any: ...
  173. def ravel(a: ArrayLike, order: _OrderKACF = ...) -> ndarray: ...
  174. def nonzero(a: ArrayLike) -> Tuple[ndarray, ...]: ...
  175. def shape(a: ArrayLike) -> _Shape: ...
  176. def compress(
  177. condition: ArrayLike, # 1D bool array
  178. a: ArrayLike,
  179. axis: Optional[int] = ...,
  180. out: Optional[ndarray] = ...,
  181. ) -> ndarray: ...
  182. @overload
  183. def clip(
  184. a: ArrayLike,
  185. a_min: ArrayLike,
  186. a_max: Optional[ArrayLike],
  187. out: Optional[ndarray] = ...,
  188. **kwargs: Any,
  189. ) -> Any: ...
  190. @overload
  191. def clip(
  192. a: ArrayLike,
  193. a_min: None,
  194. a_max: ArrayLike,
  195. out: Optional[ndarray] = ...,
  196. **kwargs: Any,
  197. ) -> Any: ...
  198. def sum(
  199. a: ArrayLike,
  200. axis: _ShapeLike = ...,
  201. dtype: DTypeLike = ...,
  202. out: Optional[ndarray] = ...,
  203. keepdims: bool = ...,
  204. initial: _NumberLike = ...,
  205. where: _ArrayLikeBool = ...,
  206. ) -> Any: ...
  207. @overload
  208. def all(
  209. a: ArrayLike,
  210. axis: None = ...,
  211. out: None = ...,
  212. keepdims: Literal[False] = ...,
  213. ) -> bool_: ...
  214. @overload
  215. def all(
  216. a: ArrayLike,
  217. axis: Optional[_ShapeLike] = ...,
  218. out: Optional[ndarray] = ...,
  219. keepdims: bool = ...,
  220. ) -> Any: ...
  221. @overload
  222. def any(
  223. a: ArrayLike,
  224. axis: None = ...,
  225. out: None = ...,
  226. keepdims: Literal[False] = ...,
  227. ) -> bool_: ...
  228. @overload
  229. def any(
  230. a: ArrayLike,
  231. axis: Optional[_ShapeLike] = ...,
  232. out: Optional[ndarray] = ...,
  233. keepdims: bool = ...,
  234. ) -> Any: ...
  235. def cumsum(
  236. a: ArrayLike,
  237. axis: Optional[int] = ...,
  238. dtype: DTypeLike = ...,
  239. out: Optional[ndarray] = ...,
  240. ) -> ndarray: ...
  241. def ptp(
  242. a: ArrayLike,
  243. axis: Optional[_ShapeLike] = ...,
  244. out: Optional[ndarray] = ...,
  245. keepdims: bool = ...,
  246. ) -> Any: ...
  247. def amax(
  248. a: ArrayLike,
  249. axis: Optional[_ShapeLike] = ...,
  250. out: Optional[ndarray] = ...,
  251. keepdims: bool = ...,
  252. initial: _NumberLike = ...,
  253. where: _ArrayLikeBool = ...,
  254. ) -> Any: ...
  255. def amin(
  256. a: ArrayLike,
  257. axis: Optional[_ShapeLike] = ...,
  258. out: Optional[ndarray] = ...,
  259. keepdims: bool = ...,
  260. initial: _NumberLike = ...,
  261. where: _ArrayLikeBool = ...,
  262. ) -> Any: ...
  263. # TODO: `np.prod()``: For object arrays `initial` does not necessarily
  264. # have to be a numerical scalar.
  265. # The only requirement is that it is compatible
  266. # with the `.__mul__()` method(s) of the passed array's elements.
  267. # Note that the same situation holds for all wrappers around
  268. # `np.ufunc.reduce`, e.g. `np.sum()` (`.__add__()`).
  269. def prod(
  270. a: ArrayLike,
  271. axis: Optional[_ShapeLike] = ...,
  272. dtype: DTypeLike = ...,
  273. out: Optional[ndarray] = ...,
  274. keepdims: bool = ...,
  275. initial: _NumberLike = ...,
  276. where: _ArrayLikeBool = ...,
  277. ) -> Any: ...
  278. def cumprod(
  279. a: ArrayLike,
  280. axis: Optional[int] = ...,
  281. dtype: DTypeLike = ...,
  282. out: Optional[ndarray] = ...,
  283. ) -> ndarray: ...
  284. def ndim(a: ArrayLike) -> int: ...
  285. def size(a: ArrayLike, axis: Optional[int] = ...) -> int: ...
  286. def around(
  287. a: ArrayLike,
  288. decimals: int = ...,
  289. out: Optional[ndarray] = ...,
  290. ) -> Any: ...
  291. def mean(
  292. a: ArrayLike,
  293. axis: Optional[_ShapeLike] = ...,
  294. dtype: DTypeLike = ...,
  295. out: Optional[ndarray] = ...,
  296. keepdims: bool = ...,
  297. ) -> Any: ...
  298. def std(
  299. a: ArrayLike,
  300. axis: Optional[_ShapeLike] = ...,
  301. dtype: DTypeLike = ...,
  302. out: Optional[ndarray] = ...,
  303. ddof: int = ...,
  304. keepdims: bool = ...,
  305. ) -> Any: ...
  306. def var(
  307. a: ArrayLike,
  308. axis: Optional[_ShapeLike] = ...,
  309. dtype: DTypeLike = ...,
  310. out: Optional[ndarray] = ...,
  311. ddof: int = ...,
  312. keepdims: bool = ...,
  313. ) -> Any: ...