fromnumeric.pyi 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343
  1. """Tests for :mod:`_core.fromnumeric`."""
  2. from typing import Any, Literal as L
  3. import numpy as np
  4. import numpy.typing as npt
  5. from typing_extensions import assert_type
  6. class NDArraySubclass(npt.NDArray[np.complex128]):
  7. ...
  8. AR_b: npt.NDArray[np.bool]
  9. AR_f4: npt.NDArray[np.float32]
  10. AR_c16: npt.NDArray[np.complex128]
  11. AR_u8: npt.NDArray[np.uint64]
  12. AR_i8: npt.NDArray[np.int64]
  13. AR_O: npt.NDArray[np.object_]
  14. AR_subclass: NDArraySubclass
  15. AR_m: npt.NDArray[np.timedelta64]
  16. AR_0d: np.ndarray[tuple[()], np.dtype[Any]]
  17. AR_1d: np.ndarray[tuple[int], np.dtype[Any]]
  18. AR_nd: np.ndarray[tuple[int, ...], np.dtype[Any]]
  19. b: np.bool
  20. f4: np.float32
  21. i8: np.int64
  22. f: float
  23. assert_type(np.take(b, 0), np.bool)
  24. assert_type(np.take(f4, 0), np.float32)
  25. assert_type(np.take(f, 0), Any)
  26. assert_type(np.take(AR_b, 0), np.bool)
  27. assert_type(np.take(AR_f4, 0), np.float32)
  28. assert_type(np.take(AR_b, [0]), npt.NDArray[np.bool])
  29. assert_type(np.take(AR_f4, [0]), npt.NDArray[np.float32])
  30. assert_type(np.take([1], [0]), npt.NDArray[Any])
  31. assert_type(np.take(AR_f4, [0], out=AR_subclass), NDArraySubclass)
  32. assert_type(np.reshape(b, 1), np.ndarray[tuple[int], np.dtype[np.bool]])
  33. assert_type(np.reshape(f4, 1), np.ndarray[tuple[int], np.dtype[np.float32]])
  34. assert_type(np.reshape(f, 1), np.ndarray[tuple[int], np.dtype[Any]])
  35. assert_type(np.reshape(AR_b, 1), np.ndarray[tuple[int], np.dtype[np.bool]])
  36. assert_type(np.reshape(AR_f4, 1), np.ndarray[tuple[int], np.dtype[np.float32]])
  37. assert_type(np.choose(1, [True, True]), Any)
  38. assert_type(np.choose([1], [True, True]), npt.NDArray[Any])
  39. assert_type(np.choose([1], AR_b), npt.NDArray[np.bool])
  40. assert_type(np.choose([1], AR_b, out=AR_f4), npt.NDArray[np.float32])
  41. assert_type(np.repeat(b, 1), npt.NDArray[np.bool])
  42. assert_type(np.repeat(f4, 1), npt.NDArray[np.float32])
  43. assert_type(np.repeat(f, 1), npt.NDArray[Any])
  44. assert_type(np.repeat(AR_b, 1), npt.NDArray[np.bool])
  45. assert_type(np.repeat(AR_f4, 1), npt.NDArray[np.float32])
  46. # TODO: array_bdd tests for np.put()
  47. assert_type(np.swapaxes([[0, 1]], 0, 0), npt.NDArray[Any])
  48. assert_type(np.swapaxes(AR_b, 0, 0), npt.NDArray[np.bool])
  49. assert_type(np.swapaxes(AR_f4, 0, 0), npt.NDArray[np.float32])
  50. assert_type(np.transpose(b), npt.NDArray[np.bool])
  51. assert_type(np.transpose(f4), npt.NDArray[np.float32])
  52. assert_type(np.transpose(f), npt.NDArray[Any])
  53. assert_type(np.transpose(AR_b), npt.NDArray[np.bool])
  54. assert_type(np.transpose(AR_f4), npt.NDArray[np.float32])
  55. assert_type(np.partition(b, 0, axis=None), npt.NDArray[np.bool])
  56. assert_type(np.partition(f4, 0, axis=None), npt.NDArray[np.float32])
  57. assert_type(np.partition(f, 0, axis=None), npt.NDArray[Any])
  58. assert_type(np.partition(AR_b, 0), npt.NDArray[np.bool])
  59. assert_type(np.partition(AR_f4, 0), npt.NDArray[np.float32])
  60. assert_type(np.argpartition(b, 0), npt.NDArray[np.intp])
  61. assert_type(np.argpartition(f4, 0), npt.NDArray[np.intp])
  62. assert_type(np.argpartition(f, 0), npt.NDArray[np.intp])
  63. assert_type(np.argpartition(AR_b, 0), npt.NDArray[np.intp])
  64. assert_type(np.argpartition(AR_f4, 0), npt.NDArray[np.intp])
  65. assert_type(np.sort([2, 1], 0), npt.NDArray[Any])
  66. assert_type(np.sort(AR_b, 0), npt.NDArray[np.bool])
  67. assert_type(np.sort(AR_f4, 0), npt.NDArray[np.float32])
  68. assert_type(np.argsort(AR_b, 0), npt.NDArray[np.intp])
  69. assert_type(np.argsort(AR_f4, 0), npt.NDArray[np.intp])
  70. assert_type(np.argmax(AR_b), np.intp)
  71. assert_type(np.argmax(AR_f4), np.intp)
  72. assert_type(np.argmax(AR_b, axis=0), Any)
  73. assert_type(np.argmax(AR_f4, axis=0), Any)
  74. assert_type(np.argmax(AR_f4, out=AR_subclass), NDArraySubclass)
  75. assert_type(np.argmin(AR_b), np.intp)
  76. assert_type(np.argmin(AR_f4), np.intp)
  77. assert_type(np.argmin(AR_b, axis=0), Any)
  78. assert_type(np.argmin(AR_f4, axis=0), Any)
  79. assert_type(np.argmin(AR_f4, out=AR_subclass), NDArraySubclass)
  80. assert_type(np.searchsorted(AR_b[0], 0), np.intp)
  81. assert_type(np.searchsorted(AR_f4[0], 0), np.intp)
  82. assert_type(np.searchsorted(AR_b[0], [0]), npt.NDArray[np.intp])
  83. assert_type(np.searchsorted(AR_f4[0], [0]), npt.NDArray[np.intp])
  84. assert_type(np.resize(b, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.bool]])
  85. assert_type(np.resize(f4, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.float32]])
  86. assert_type(np.resize(f, (5, 5)), np.ndarray[tuple[int, int], np.dtype[Any]])
  87. assert_type(np.resize(AR_b, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.bool]])
  88. assert_type(np.resize(AR_f4, (5, 5)), np.ndarray[tuple[int, int], np.dtype[np.float32]])
  89. assert_type(np.squeeze(b), np.bool)
  90. assert_type(np.squeeze(f4), np.float32)
  91. assert_type(np.squeeze(f), npt.NDArray[Any])
  92. assert_type(np.squeeze(AR_b), npt.NDArray[np.bool])
  93. assert_type(np.squeeze(AR_f4), npt.NDArray[np.float32])
  94. assert_type(np.diagonal(AR_b), npt.NDArray[np.bool])
  95. assert_type(np.diagonal(AR_f4), npt.NDArray[np.float32])
  96. assert_type(np.trace(AR_b), Any)
  97. assert_type(np.trace(AR_f4), Any)
  98. assert_type(np.trace(AR_f4, out=AR_subclass), NDArraySubclass)
  99. assert_type(np.ravel(b), np.ndarray[tuple[int], np.dtype[np.bool]])
  100. assert_type(np.ravel(f4), np.ndarray[tuple[int], np.dtype[np.float32]])
  101. assert_type(np.ravel(f), np.ndarray[tuple[int], np.dtype[np.float64 | np.int_ | np.bool]])
  102. assert_type(np.ravel(AR_b), np.ndarray[tuple[int], np.dtype[np.bool]])
  103. assert_type(np.ravel(AR_f4), np.ndarray[tuple[int], np.dtype[np.float32]])
  104. assert_type(np.nonzero(AR_b), tuple[npt.NDArray[np.intp], ...])
  105. assert_type(np.nonzero(AR_f4), tuple[npt.NDArray[np.intp], ...])
  106. assert_type(np.nonzero(AR_1d), tuple[npt.NDArray[np.intp], ...])
  107. assert_type(np.nonzero(AR_nd), tuple[npt.NDArray[np.intp], ...])
  108. assert_type(np.shape(b), tuple[()])
  109. assert_type(np.shape(f), tuple[()])
  110. assert_type(np.shape([1]), tuple[int])
  111. assert_type(np.shape([[2]]), tuple[int, int])
  112. assert_type(np.shape([[[3]]]), tuple[int, ...])
  113. assert_type(np.shape(AR_b), tuple[int, ...])
  114. assert_type(np.shape(AR_nd), tuple[int, ...])
  115. # these fail on mypy, but it works as expected with pyright/pylance
  116. # assert_type(np.shape(AR_0d), tuple[()])
  117. # assert_type(np.shape(AR_1d), tuple[int])
  118. # assert_type(np.shape(AR_2d), tuple[int, int])
  119. assert_type(np.compress([True], b), npt.NDArray[np.bool])
  120. assert_type(np.compress([True], f4), npt.NDArray[np.float32])
  121. assert_type(np.compress([True], f), npt.NDArray[Any])
  122. assert_type(np.compress([True], AR_b), npt.NDArray[np.bool])
  123. assert_type(np.compress([True], AR_f4), npt.NDArray[np.float32])
  124. assert_type(np.clip(b, 0, 1.0), np.bool)
  125. assert_type(np.clip(f4, -1, 1), np.float32)
  126. assert_type(np.clip(f, 0, 1), Any)
  127. assert_type(np.clip(AR_b, 0, 1), npt.NDArray[np.bool])
  128. assert_type(np.clip(AR_f4, 0, 1), npt.NDArray[np.float32])
  129. assert_type(np.clip([0], 0, 1), npt.NDArray[Any])
  130. assert_type(np.clip(AR_b, 0, 1, out=AR_subclass), NDArraySubclass)
  131. assert_type(np.sum(b), np.bool)
  132. assert_type(np.sum(f4), np.float32)
  133. assert_type(np.sum(f), Any)
  134. assert_type(np.sum(AR_b), np.bool)
  135. assert_type(np.sum(AR_f4), np.float32)
  136. assert_type(np.sum(AR_b, axis=0), Any)
  137. assert_type(np.sum(AR_f4, axis=0), Any)
  138. assert_type(np.sum(AR_f4, out=AR_subclass), NDArraySubclass)
  139. assert_type(np.sum(AR_f4, dtype=np.float64), np.float64)
  140. assert_type(np.sum(AR_f4, None, np.float64), np.float64)
  141. assert_type(np.sum(AR_f4, dtype=np.float64, keepdims=False), np.float64)
  142. assert_type(np.sum(AR_f4, None, np.float64, keepdims=False), np.float64)
  143. assert_type(np.sum(AR_f4, dtype=np.float64, keepdims=True), np.float64 | npt.NDArray[np.float64])
  144. assert_type(np.sum(AR_f4, None, np.float64, keepdims=True), np.float64 | npt.NDArray[np.float64])
  145. assert_type(np.all(b), np.bool)
  146. assert_type(np.all(f4), np.bool)
  147. assert_type(np.all(f), np.bool)
  148. assert_type(np.all(AR_b), np.bool)
  149. assert_type(np.all(AR_f4), np.bool)
  150. assert_type(np.all(AR_b, axis=0), Any)
  151. assert_type(np.all(AR_f4, axis=0), Any)
  152. assert_type(np.all(AR_b, keepdims=True), Any)
  153. assert_type(np.all(AR_f4, keepdims=True), Any)
  154. assert_type(np.all(AR_f4, out=AR_subclass), NDArraySubclass)
  155. assert_type(np.any(b), np.bool)
  156. assert_type(np.any(f4), np.bool)
  157. assert_type(np.any(f), np.bool)
  158. assert_type(np.any(AR_b), np.bool)
  159. assert_type(np.any(AR_f4), np.bool)
  160. assert_type(np.any(AR_b, axis=0), Any)
  161. assert_type(np.any(AR_f4, axis=0), Any)
  162. assert_type(np.any(AR_b, keepdims=True), Any)
  163. assert_type(np.any(AR_f4, keepdims=True), Any)
  164. assert_type(np.any(AR_f4, out=AR_subclass), NDArraySubclass)
  165. assert_type(np.cumsum(b), npt.NDArray[np.bool])
  166. assert_type(np.cumsum(f4), npt.NDArray[np.float32])
  167. assert_type(np.cumsum(f), npt.NDArray[Any])
  168. assert_type(np.cumsum(AR_b), npt.NDArray[np.bool])
  169. assert_type(np.cumsum(AR_f4), npt.NDArray[np.float32])
  170. assert_type(np.cumsum(f, dtype=float), npt.NDArray[Any])
  171. assert_type(np.cumsum(f, dtype=np.float64), npt.NDArray[np.float64])
  172. assert_type(np.cumsum(AR_f4, out=AR_subclass), NDArraySubclass)
  173. assert_type(np.cumulative_sum(b), npt.NDArray[np.bool])
  174. assert_type(np.cumulative_sum(f4), npt.NDArray[np.float32])
  175. assert_type(np.cumulative_sum(f), npt.NDArray[Any])
  176. assert_type(np.cumulative_sum(AR_b), npt.NDArray[np.bool])
  177. assert_type(np.cumulative_sum(AR_f4), npt.NDArray[np.float32])
  178. assert_type(np.cumulative_sum(f, dtype=float), npt.NDArray[Any])
  179. assert_type(np.cumulative_sum(f, dtype=np.float64), npt.NDArray[np.float64])
  180. assert_type(np.cumulative_sum(AR_f4, out=AR_subclass), NDArraySubclass)
  181. assert_type(np.ptp(b), np.bool)
  182. assert_type(np.ptp(f4), np.float32)
  183. assert_type(np.ptp(f), Any)
  184. assert_type(np.ptp(AR_b), np.bool)
  185. assert_type(np.ptp(AR_f4), np.float32)
  186. assert_type(np.ptp(AR_b, axis=0), Any)
  187. assert_type(np.ptp(AR_f4, axis=0), Any)
  188. assert_type(np.ptp(AR_b, keepdims=True), Any)
  189. assert_type(np.ptp(AR_f4, keepdims=True), Any)
  190. assert_type(np.ptp(AR_f4, out=AR_subclass), NDArraySubclass)
  191. assert_type(np.amax(b), np.bool)
  192. assert_type(np.amax(f4), np.float32)
  193. assert_type(np.amax(f), Any)
  194. assert_type(np.amax(AR_b), np.bool)
  195. assert_type(np.amax(AR_f4), np.float32)
  196. assert_type(np.amax(AR_b, axis=0), Any)
  197. assert_type(np.amax(AR_f4, axis=0), Any)
  198. assert_type(np.amax(AR_b, keepdims=True), Any)
  199. assert_type(np.amax(AR_f4, keepdims=True), Any)
  200. assert_type(np.amax(AR_f4, out=AR_subclass), NDArraySubclass)
  201. assert_type(np.amin(b), np.bool)
  202. assert_type(np.amin(f4), np.float32)
  203. assert_type(np.amin(f), Any)
  204. assert_type(np.amin(AR_b), np.bool)
  205. assert_type(np.amin(AR_f4), np.float32)
  206. assert_type(np.amin(AR_b, axis=0), Any)
  207. assert_type(np.amin(AR_f4, axis=0), Any)
  208. assert_type(np.amin(AR_b, keepdims=True), Any)
  209. assert_type(np.amin(AR_f4, keepdims=True), Any)
  210. assert_type(np.amin(AR_f4, out=AR_subclass), NDArraySubclass)
  211. assert_type(np.prod(AR_b), np.int_)
  212. assert_type(np.prod(AR_u8), np.uint64)
  213. assert_type(np.prod(AR_i8), np.int64)
  214. assert_type(np.prod(AR_f4), np.floating[Any])
  215. assert_type(np.prod(AR_c16), np.complexfloating[Any, Any])
  216. assert_type(np.prod(AR_O), Any)
  217. assert_type(np.prod(AR_f4, axis=0), Any)
  218. assert_type(np.prod(AR_f4, keepdims=True), Any)
  219. assert_type(np.prod(AR_f4, dtype=np.float64), np.float64)
  220. assert_type(np.prod(AR_f4, dtype=float), Any)
  221. assert_type(np.prod(AR_f4, out=AR_subclass), NDArraySubclass)
  222. assert_type(np.cumprod(AR_b), npt.NDArray[np.int_])
  223. assert_type(np.cumprod(AR_u8), npt.NDArray[np.uint64])
  224. assert_type(np.cumprod(AR_i8), npt.NDArray[np.int64])
  225. assert_type(np.cumprod(AR_f4), npt.NDArray[np.floating[Any]])
  226. assert_type(np.cumprod(AR_c16), npt.NDArray[np.complexfloating[Any, Any]])
  227. assert_type(np.cumprod(AR_O), npt.NDArray[np.object_])
  228. assert_type(np.cumprod(AR_f4, axis=0), npt.NDArray[np.floating[Any]])
  229. assert_type(np.cumprod(AR_f4, dtype=np.float64), npt.NDArray[np.float64])
  230. assert_type(np.cumprod(AR_f4, dtype=float), npt.NDArray[Any])
  231. assert_type(np.cumprod(AR_f4, out=AR_subclass), NDArraySubclass)
  232. assert_type(np.cumulative_prod(AR_b), npt.NDArray[np.int_])
  233. assert_type(np.cumulative_prod(AR_u8), npt.NDArray[np.uint64])
  234. assert_type(np.cumulative_prod(AR_i8), npt.NDArray[np.int64])
  235. assert_type(np.cumulative_prod(AR_f4), npt.NDArray[np.floating[Any]])
  236. assert_type(np.cumulative_prod(AR_c16), npt.NDArray[np.complexfloating[Any, Any]])
  237. assert_type(np.cumulative_prod(AR_O), npt.NDArray[np.object_])
  238. assert_type(np.cumulative_prod(AR_f4, axis=0), npt.NDArray[np.floating[Any]])
  239. assert_type(np.cumulative_prod(AR_f4, dtype=np.float64), npt.NDArray[np.float64])
  240. assert_type(np.cumulative_prod(AR_f4, dtype=float), npt.NDArray[Any])
  241. assert_type(np.cumulative_prod(AR_f4, out=AR_subclass), NDArraySubclass)
  242. assert_type(np.ndim(b), int)
  243. assert_type(np.ndim(f4), int)
  244. assert_type(np.ndim(f), int)
  245. assert_type(np.ndim(AR_b), int)
  246. assert_type(np.ndim(AR_f4), int)
  247. assert_type(np.size(b), int)
  248. assert_type(np.size(f4), int)
  249. assert_type(np.size(f), int)
  250. assert_type(np.size(AR_b), int)
  251. assert_type(np.size(AR_f4), int)
  252. assert_type(np.around(b), np.float16)
  253. assert_type(np.around(f), Any)
  254. assert_type(np.around(i8), np.int64)
  255. assert_type(np.around(f4), np.float32)
  256. assert_type(np.around(AR_b), npt.NDArray[np.float16])
  257. assert_type(np.around(AR_i8), npt.NDArray[np.int64])
  258. assert_type(np.around(AR_f4), npt.NDArray[np.float32])
  259. assert_type(np.around([1.5]), npt.NDArray[Any])
  260. assert_type(np.around(AR_f4, out=AR_subclass), NDArraySubclass)
  261. assert_type(np.mean(AR_b), np.floating[Any])
  262. assert_type(np.mean(AR_i8), np.floating[Any])
  263. assert_type(np.mean(AR_f4), np.floating[Any])
  264. assert_type(np.mean(AR_m), np.timedelta64)
  265. assert_type(np.mean(AR_c16), np.complexfloating[Any, Any])
  266. assert_type(np.mean(AR_O), Any)
  267. assert_type(np.mean(AR_f4, axis=0), Any)
  268. assert_type(np.mean(AR_f4, keepdims=True), Any)
  269. assert_type(np.mean(AR_f4, dtype=float), Any)
  270. assert_type(np.mean(AR_f4, dtype=np.float64), np.float64)
  271. assert_type(np.mean(AR_f4, out=AR_subclass), NDArraySubclass)
  272. assert_type(np.mean(AR_f4, dtype=np.float64), np.float64)
  273. assert_type(np.mean(AR_f4, None, np.float64), np.float64)
  274. assert_type(np.mean(AR_f4, dtype=np.float64, keepdims=False), np.float64)
  275. assert_type(np.mean(AR_f4, None, np.float64, keepdims=False), np.float64)
  276. assert_type(np.mean(AR_f4, dtype=np.float64, keepdims=True), np.float64 | npt.NDArray[np.float64])
  277. assert_type(np.mean(AR_f4, None, np.float64, keepdims=True), np.float64 | npt.NDArray[np.float64])
  278. assert_type(np.std(AR_b), np.floating[Any])
  279. assert_type(np.std(AR_i8), np.floating[Any])
  280. assert_type(np.std(AR_f4), np.floating[Any])
  281. assert_type(np.std(AR_c16), np.floating[Any])
  282. assert_type(np.std(AR_O), Any)
  283. assert_type(np.std(AR_f4, axis=0), Any)
  284. assert_type(np.std(AR_f4, keepdims=True), Any)
  285. assert_type(np.std(AR_f4, dtype=float), Any)
  286. assert_type(np.std(AR_f4, dtype=np.float64), np.float64)
  287. assert_type(np.std(AR_f4, out=AR_subclass), NDArraySubclass)
  288. assert_type(np.var(AR_b), np.floating[Any])
  289. assert_type(np.var(AR_i8), np.floating[Any])
  290. assert_type(np.var(AR_f4), np.floating[Any])
  291. assert_type(np.var(AR_c16), np.floating[Any])
  292. assert_type(np.var(AR_O), Any)
  293. assert_type(np.var(AR_f4, axis=0), Any)
  294. assert_type(np.var(AR_f4, keepdims=True), Any)
  295. assert_type(np.var(AR_f4, dtype=float), Any)
  296. assert_type(np.var(AR_f4, dtype=np.float64), np.float64)
  297. assert_type(np.var(AR_f4, out=AR_subclass), NDArraySubclass)