multiarray.pyi 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196
  1. import datetime as dt
  2. from typing import Any, Literal, TypeVar
  3. import numpy as np
  4. import numpy.typing as npt
  5. from typing_extensions import Unpack, assert_type
  6. _SCT = TypeVar("_SCT", bound=np.generic, covariant=True)
  7. class SubClass(npt.NDArray[_SCT]): ...
  8. subclass: SubClass[np.float64]
  9. AR_f8: npt.NDArray[np.float64]
  10. AR_i8: npt.NDArray[np.int64]
  11. AR_u1: npt.NDArray[np.uint8]
  12. AR_m: npt.NDArray[np.timedelta64]
  13. AR_M: npt.NDArray[np.datetime64]
  14. AR_LIKE_f: list[float]
  15. AR_LIKE_i: list[int]
  16. m: np.timedelta64
  17. M: np.datetime64
  18. b_f8 = np.broadcast(AR_f8)
  19. b_i8_f8_f8 = np.broadcast(AR_i8, AR_f8, AR_f8)
  20. nditer_obj: np.nditer
  21. date_scalar: dt.date
  22. date_seq: list[dt.date]
  23. timedelta_seq: list[dt.timedelta]
  24. n1: Literal[1]
  25. n2: Literal[2]
  26. n3: Literal[3]
  27. f8: np.float64
  28. def func11(a: int) -> bool: ...
  29. def func21(a: int, b: int) -> int: ...
  30. def func12(a: int) -> tuple[complex, bool]: ...
  31. assert_type(next(b_f8), tuple[Any, ...])
  32. assert_type(b_f8.reset(), None)
  33. assert_type(b_f8.index, int)
  34. assert_type(b_f8.iters, tuple[np.flatiter[Any], ...])
  35. assert_type(b_f8.nd, int)
  36. assert_type(b_f8.ndim, int)
  37. assert_type(b_f8.numiter, int)
  38. assert_type(b_f8.shape, tuple[int, ...])
  39. assert_type(b_f8.size, int)
  40. assert_type(next(b_i8_f8_f8), tuple[Any, ...])
  41. assert_type(b_i8_f8_f8.reset(), None)
  42. assert_type(b_i8_f8_f8.index, int)
  43. assert_type(b_i8_f8_f8.iters, tuple[np.flatiter[Any], ...])
  44. assert_type(b_i8_f8_f8.nd, int)
  45. assert_type(b_i8_f8_f8.ndim, int)
  46. assert_type(b_i8_f8_f8.numiter, int)
  47. assert_type(b_i8_f8_f8.shape, tuple[int, ...])
  48. assert_type(b_i8_f8_f8.size, int)
  49. assert_type(np.inner(AR_f8, AR_i8), Any)
  50. assert_type(np.where([True, True, False]), tuple[npt.NDArray[np.intp], ...])
  51. assert_type(np.where([True, True, False], 1, 0), npt.NDArray[Any])
  52. assert_type(np.lexsort([0, 1, 2]), Any)
  53. assert_type(np.can_cast(np.dtype("i8"), int), bool)
  54. assert_type(np.can_cast(AR_f8, "f8"), bool)
  55. assert_type(np.can_cast(AR_f8, np.complex128, casting="unsafe"), bool)
  56. assert_type(np.min_scalar_type([1]), np.dtype[Any])
  57. assert_type(np.min_scalar_type(AR_f8), np.dtype[Any])
  58. assert_type(np.result_type(int, [1]), np.dtype[Any])
  59. assert_type(np.result_type(AR_f8, AR_u1), np.dtype[Any])
  60. assert_type(np.result_type(AR_f8, np.complex128), np.dtype[Any])
  61. assert_type(np.dot(AR_LIKE_f, AR_i8), Any)
  62. assert_type(np.dot(AR_u1, 1), Any)
  63. assert_type(np.dot(1.5j, 1), Any)
  64. assert_type(np.dot(AR_u1, 1, out=AR_f8), npt.NDArray[np.float64])
  65. assert_type(np.vdot(AR_LIKE_f, AR_i8), np.floating[Any])
  66. assert_type(np.vdot(AR_u1, 1), np.signedinteger[Any])
  67. assert_type(np.vdot(1.5j, 1), np.complexfloating[Any, Any])
  68. assert_type(np.bincount(AR_i8), npt.NDArray[np.intp])
  69. assert_type(np.copyto(AR_f8, [1., 1.5, 1.6]), None)
  70. assert_type(np.putmask(AR_f8, [True, True, False], 1.5), None)
  71. assert_type(np.packbits(AR_i8), npt.NDArray[np.uint8])
  72. assert_type(np.packbits(AR_u1), npt.NDArray[np.uint8])
  73. assert_type(np.unpackbits(AR_u1), npt.NDArray[np.uint8])
  74. assert_type(np.shares_memory(1, 2), bool)
  75. assert_type(np.shares_memory(AR_f8, AR_f8, max_work=1), bool)
  76. assert_type(np.may_share_memory(1, 2), bool)
  77. assert_type(np.may_share_memory(AR_f8, AR_f8, max_work=1), bool)
  78. assert_type(np.promote_types(np.int32, np.int64), np.dtype[Any])
  79. assert_type(np.promote_types("f4", float), np.dtype[Any])
  80. assert_type(np.frompyfunc(func11, n1, n1).nin, Literal[1])
  81. assert_type(np.frompyfunc(func11, n1, n1).nout, Literal[1])
  82. assert_type(np.frompyfunc(func11, n1, n1).nargs, Literal[2])
  83. assert_type(np.frompyfunc(func11, n1, n1).ntypes, Literal[1])
  84. assert_type(np.frompyfunc(func11, n1, n1).identity, None)
  85. assert_type(np.frompyfunc(func11, n1, n1).signature, None)
  86. assert_type(np.frompyfunc(func11, n1, n1)(f8), bool)
  87. assert_type(np.frompyfunc(func11, n1, n1)(AR_f8), bool | npt.NDArray[np.object_])
  88. assert_type(np.frompyfunc(func11, n1, n1).at(AR_f8, AR_i8), None)
  89. assert_type(np.frompyfunc(func21, n2, n1).nin, Literal[2])
  90. assert_type(np.frompyfunc(func21, n2, n1).nout, Literal[1])
  91. assert_type(np.frompyfunc(func21, n2, n1).nargs, Literal[3])
  92. assert_type(np.frompyfunc(func21, n2, n1).ntypes, Literal[1])
  93. assert_type(np.frompyfunc(func21, n2, n1).identity, None)
  94. assert_type(np.frompyfunc(func21, n2, n1).signature, None)
  95. assert_type(np.frompyfunc(func21, n2, n1)(f8, f8), int)
  96. assert_type(np.frompyfunc(func21, n2, n1)(AR_f8, f8), int | npt.NDArray[np.object_])
  97. assert_type(np.frompyfunc(func21, n2, n1)(f8, AR_f8), int | npt.NDArray[np.object_])
  98. assert_type(np.frompyfunc(func21, n2, n1).reduce(AR_f8, axis=0), int | npt.NDArray[np.object_])
  99. assert_type(np.frompyfunc(func21, n2, n1).accumulate(AR_f8), npt.NDArray[np.object_])
  100. assert_type(np.frompyfunc(func21, n2, n1).reduceat(AR_f8, AR_i8), npt.NDArray[np.object_])
  101. assert_type(np.frompyfunc(func21, n2, n1).outer(f8, f8), int)
  102. assert_type(np.frompyfunc(func21, n2, n1).outer(AR_f8, f8), int | npt.NDArray[np.object_])
  103. assert_type(np.frompyfunc(func21, n2, n1, identity=0).nin, Literal[2])
  104. assert_type(np.frompyfunc(func21, n2, n1, identity=0).nout, Literal[1])
  105. assert_type(np.frompyfunc(func21, n2, n1, identity=0).nargs, Literal[3])
  106. assert_type(np.frompyfunc(func21, n2, n1, identity=0).ntypes, Literal[1])
  107. assert_type(np.frompyfunc(func21, n2, n1, identity=0).identity, int)
  108. assert_type(np.frompyfunc(func21, n2, n1, identity=0).signature, None)
  109. assert_type(np.frompyfunc(func12, n1, n2).nin, Literal[1])
  110. assert_type(np.frompyfunc(func12, n1, n2).nout, Literal[2])
  111. assert_type(np.frompyfunc(func12, n1, n2).nargs, int)
  112. assert_type(np.frompyfunc(func12, n1, n2).ntypes, Literal[1])
  113. assert_type(np.frompyfunc(func12, n1, n2).identity, None)
  114. assert_type(np.frompyfunc(func12, n1, n2).signature, None)
  115. assert_type(
  116. np.frompyfunc(func12, n2, n2)(f8, f8),
  117. tuple[complex, complex, Unpack[tuple[complex, ...]]],
  118. )
  119. assert_type(
  120. np.frompyfunc(func12, n2, n2)(AR_f8, f8),
  121. tuple[
  122. complex | npt.NDArray[np.object_],
  123. complex | npt.NDArray[np.object_],
  124. Unpack[tuple[complex | npt.NDArray[np.object_], ...]],
  125. ],
  126. )
  127. assert_type(np.datetime_data("m8[D]"), tuple[str, int])
  128. assert_type(np.datetime_data(np.datetime64), tuple[str, int])
  129. assert_type(np.datetime_data(np.dtype(np.timedelta64)), tuple[str, int])
  130. assert_type(np.busday_count("2011-01", "2011-02"), np.int_)
  131. assert_type(np.busday_count(["2011-01"], "2011-02"), npt.NDArray[np.int_])
  132. assert_type(np.busday_count(["2011-01"], date_scalar), npt.NDArray[np.int_])
  133. assert_type(np.busday_offset(M, m), np.datetime64)
  134. assert_type(np.busday_offset(date_scalar, m), np.datetime64)
  135. assert_type(np.busday_offset(M, 5), np.datetime64)
  136. assert_type(np.busday_offset(AR_M, m), npt.NDArray[np.datetime64])
  137. assert_type(np.busday_offset(M, timedelta_seq), npt.NDArray[np.datetime64])
  138. assert_type(np.busday_offset("2011-01", "2011-02", roll="forward"), np.datetime64)
  139. assert_type(np.busday_offset(["2011-01"], "2011-02", roll="forward"), npt.NDArray[np.datetime64])
  140. assert_type(np.is_busday("2012"), np.bool)
  141. assert_type(np.is_busday(date_scalar), np.bool)
  142. assert_type(np.is_busday(["2012"]), npt.NDArray[np.bool])
  143. assert_type(np.datetime_as_string(M), np.str_)
  144. assert_type(np.datetime_as_string(AR_M), npt.NDArray[np.str_])
  145. assert_type(np.busdaycalendar(holidays=date_seq), np.busdaycalendar)
  146. assert_type(np.busdaycalendar(holidays=[M]), np.busdaycalendar)
  147. assert_type(np.char.compare_chararrays("a", "b", "!=", rstrip=False), npt.NDArray[np.bool])
  148. assert_type(np.char.compare_chararrays(b"a", b"a", "==", True), npt.NDArray[np.bool])
  149. assert_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], flags=["c_index"]), tuple[np.nditer, ...])
  150. assert_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_flags=[["readonly", "readonly"]]), tuple[np.nditer, ...])
  151. assert_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], op_dtypes=np.int_), tuple[np.nditer, ...])
  152. assert_type(np.nested_iters([AR_i8, AR_i8], [[0], [1]], order="C", casting="no"), tuple[np.nditer, ...])