mod.pyi 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181
  1. import datetime as dt
  2. from typing import Literal as L
  3. from typing_extensions import assert_type
  4. import numpy as np
  5. import numpy.typing as npt
  6. from numpy._typing import _64Bit
  7. f8: np.float64
  8. i8: np.int64
  9. u8: np.uint64
  10. f4: np.float32
  11. i4: np.int32
  12. u4: np.uint32
  13. m: np.timedelta64
  14. m_nat: np.timedelta64[None]
  15. m_int0: np.timedelta64[L[0]]
  16. m_int: np.timedelta64[int]
  17. m_td: np.timedelta64[dt.timedelta]
  18. b_: np.bool
  19. b: bool
  20. i: int
  21. f: float
  22. AR_b: npt.NDArray[np.bool]
  23. AR_m: npt.NDArray[np.timedelta64]
  24. # Time structures
  25. assert_type(m % m, np.timedelta64)
  26. assert_type(m % m_nat, np.timedelta64[None])
  27. assert_type(m % m_int0, np.timedelta64[None])
  28. assert_type(m % m_int, np.timedelta64[int | None])
  29. assert_type(m_nat % m, np.timedelta64[None])
  30. assert_type(m_int % m_nat, np.timedelta64[None])
  31. assert_type(m_int % m_int0, np.timedelta64[None])
  32. assert_type(m_int % m_int, np.timedelta64[int | None])
  33. assert_type(m_int % m_td, np.timedelta64[int | None])
  34. assert_type(m_td % m_nat, np.timedelta64[None])
  35. assert_type(m_td % m_int0, np.timedelta64[None])
  36. assert_type(m_td % m_int, np.timedelta64[int | None])
  37. assert_type(m_td % m_td, np.timedelta64[dt.timedelta | None])
  38. assert_type(AR_m % m, npt.NDArray[np.timedelta64])
  39. assert_type(m % AR_m, npt.NDArray[np.timedelta64])
  40. assert_type(divmod(m, m), tuple[np.int64, np.timedelta64])
  41. assert_type(divmod(m, m_nat), tuple[np.int64, np.timedelta64[None]])
  42. assert_type(divmod(m, m_int0), tuple[np.int64, np.timedelta64[None]])
  43. # workarounds for https://github.com/microsoft/pyright/issues/9663
  44. assert_type(m.__divmod__(m_int), tuple[np.int64, np.timedelta64[int | None]])
  45. assert_type(divmod(m_nat, m), tuple[np.int64, np.timedelta64[None]])
  46. assert_type(divmod(m_int, m_nat), tuple[np.int64, np.timedelta64[None]])
  47. assert_type(divmod(m_int, m_int0), tuple[np.int64, np.timedelta64[None]])
  48. assert_type(divmod(m_int, m_int), tuple[np.int64, np.timedelta64[int | None]])
  49. assert_type(divmod(m_int, m_td), tuple[np.int64, np.timedelta64[int | None]])
  50. assert_type(divmod(m_td, m_nat), tuple[np.int64, np.timedelta64[None]])
  51. assert_type(divmod(m_td, m_int0), tuple[np.int64, np.timedelta64[None]])
  52. assert_type(divmod(m_td, m_int), tuple[np.int64, np.timedelta64[int | None]])
  53. assert_type(divmod(m_td, m_td), tuple[np.int64, np.timedelta64[dt.timedelta | None]])
  54. assert_type(divmod(AR_m, m), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]])
  55. assert_type(divmod(m, AR_m), tuple[npt.NDArray[np.int64], npt.NDArray[np.timedelta64]])
  56. # Bool
  57. assert_type(b_ % b, np.int8)
  58. assert_type(b_ % i, np.int_)
  59. assert_type(b_ % f, np.float64)
  60. assert_type(b_ % b_, np.int8)
  61. assert_type(b_ % i8, np.int64)
  62. assert_type(b_ % u8, np.uint64)
  63. assert_type(b_ % f8, np.float64)
  64. assert_type(b_ % AR_b, npt.NDArray[np.int8])
  65. assert_type(divmod(b_, b), tuple[np.int8, np.int8])
  66. assert_type(divmod(b_, b_), tuple[np.int8, np.int8])
  67. # workarounds for https://github.com/microsoft/pyright/issues/9663
  68. assert_type(b_.__divmod__(i), tuple[np.int_, np.int_])
  69. assert_type(b_.__divmod__(f), tuple[np.float64, np.float64])
  70. assert_type(b_.__divmod__(i8), tuple[np.int64, np.int64])
  71. assert_type(b_.__divmod__(u8), tuple[np.uint64, np.uint64])
  72. assert_type(divmod(b_, f8), tuple[np.float64, np.float64])
  73. assert_type(divmod(b_, AR_b), tuple[npt.NDArray[np.int8], npt.NDArray[np.int8]])
  74. assert_type(b % b_, np.int8)
  75. assert_type(i % b_, np.int_)
  76. assert_type(f % b_, np.float64)
  77. assert_type(b_ % b_, np.int8)
  78. assert_type(i8 % b_, np.int64)
  79. assert_type(u8 % b_, np.uint64)
  80. assert_type(f8 % b_, np.float64)
  81. assert_type(AR_b % b_, npt.NDArray[np.int8])
  82. assert_type(divmod(b, b_), tuple[np.int8, np.int8])
  83. assert_type(divmod(i, b_), tuple[np.int_, np.int_])
  84. assert_type(divmod(f, b_), tuple[np.float64, np.float64])
  85. assert_type(divmod(b_, b_), tuple[np.int8, np.int8])
  86. assert_type(divmod(i8, b_), tuple[np.int64, np.int64])
  87. assert_type(divmod(u8, b_), tuple[np.uint64, np.uint64])
  88. assert_type(divmod(f8, b_), tuple[np.float64, np.float64])
  89. assert_type(divmod(AR_b, b_), tuple[npt.NDArray[np.int8], npt.NDArray[np.int8]])
  90. # int
  91. assert_type(i8 % b, np.int64)
  92. assert_type(i8 % i8, np.int64)
  93. assert_type(i8 % f, np.float64 | np.floating[_64Bit])
  94. assert_type(i8 % f8, np.float64 | np.floating[_64Bit])
  95. assert_type(i4 % i8, np.int64 | np.int32)
  96. assert_type(i4 % f8, np.float64 | np.float32)
  97. assert_type(i4 % i4, np.int32)
  98. assert_type(i4 % f4, np.float32)
  99. assert_type(i8 % AR_b, npt.NDArray[np.int64])
  100. assert_type(divmod(i8, b), tuple[np.int64, np.int64])
  101. assert_type(divmod(i8, i4), tuple[np.int64, np.int64] | tuple[np.int32, np.int32])
  102. assert_type(divmod(i8, i8), tuple[np.int64, np.int64])
  103. # workarounds for https://github.com/microsoft/pyright/issues/9663
  104. assert_type(i8.__divmod__(f), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64])
  105. assert_type(i8.__divmod__(f8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64])
  106. assert_type(divmod(i8, f4), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float32, np.float32])
  107. assert_type(divmod(i4, i4), tuple[np.int32, np.int32])
  108. assert_type(divmod(i4, f4), tuple[np.float32, np.float32])
  109. assert_type(divmod(i8, AR_b), tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]])
  110. assert_type(b % i8, np.int64)
  111. assert_type(f % i8, np.float64 | np.floating[_64Bit])
  112. assert_type(i8 % i8, np.int64)
  113. assert_type(f8 % i8, np.float64)
  114. assert_type(i8 % i4, np.int64 | np.int32)
  115. assert_type(f8 % i4, np.float64)
  116. assert_type(i4 % i4, np.int32)
  117. assert_type(f4 % i4, np.float32)
  118. assert_type(AR_b % i8, npt.NDArray[np.int64])
  119. assert_type(divmod(b, i8), tuple[np.int64, np.int64])
  120. assert_type(divmod(f, i8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float64, np.float64])
  121. assert_type(divmod(i8, i8), tuple[np.int64, np.int64])
  122. assert_type(divmod(f8, i8), tuple[np.float64, np.float64])
  123. assert_type(divmod(i4, i8), tuple[np.int64, np.int64] | tuple[np.int32, np.int32])
  124. assert_type(divmod(i4, i4), tuple[np.int32, np.int32])
  125. # workarounds for https://github.com/microsoft/pyright/issues/9663
  126. assert_type(f4.__divmod__(i8), tuple[np.floating[_64Bit], np.floating[_64Bit]] | tuple[np.float32, np.float32])
  127. assert_type(f4.__divmod__(i4), tuple[np.float32, np.float32])
  128. assert_type(AR_b.__divmod__(i8), tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]])
  129. # float
  130. assert_type(f8 % b, np.float64)
  131. assert_type(f8 % f, np.float64)
  132. assert_type(i8 % f4, np.floating[_64Bit] | np.float32)
  133. assert_type(f4 % f4, np.float32)
  134. assert_type(f8 % AR_b, npt.NDArray[np.float64])
  135. assert_type(divmod(f8, b), tuple[np.float64, np.float64])
  136. assert_type(divmod(f8, f), tuple[np.float64, np.float64])
  137. assert_type(divmod(f8, f8), tuple[np.float64, np.float64])
  138. assert_type(divmod(f8, f4), tuple[np.float64, np.float64])
  139. assert_type(divmod(f4, f4), tuple[np.float32, np.float32])
  140. assert_type(divmod(f8, AR_b), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])
  141. assert_type(b % f8, np.float64)
  142. assert_type(f % f8, np.float64) # pyright: ignore[reportAssertTypeFailure] # pyright incorrectly infers `builtins.float`
  143. assert_type(f8 % f8, np.float64)
  144. assert_type(f8 % f8, np.float64)
  145. assert_type(f4 % f4, np.float32)
  146. assert_type(AR_b % f8, npt.NDArray[np.float64])
  147. assert_type(divmod(b, f8), tuple[np.float64, np.float64])
  148. assert_type(divmod(f8, f8), tuple[np.float64, np.float64])
  149. assert_type(divmod(f4, f4), tuple[np.float32, np.float32])
  150. # workarounds for https://github.com/microsoft/pyright/issues/9663
  151. assert_type(f8.__rdivmod__(f), tuple[np.float64, np.float64])
  152. assert_type(f8.__rdivmod__(f4), tuple[np.float64, np.float64])
  153. assert_type(AR_b.__divmod__(f8), tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]])