scalars.pyi 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193
  1. from typing import Any, Literal, TypeAlias
  2. from typing_extensions import Unpack, assert_type
  3. import numpy as np
  4. import numpy.typing as npt
  5. _1: TypeAlias = Literal[1]
  6. b: np.bool
  7. u8: np.uint64
  8. i8: np.int64
  9. f8: np.float64
  10. c8: np.complex64
  11. c16: np.complex128
  12. m: np.timedelta64
  13. U: np.str_
  14. S: np.bytes_
  15. V: np.void
  16. O: np.object_ # cannot exists at runtime
  17. array_nd: np.ndarray[Any, Any]
  18. array_0d: np.ndarray[tuple[()], Any]
  19. array_2d_2x2: np.ndarray[tuple[Literal[2], Literal[2]], Any]
  20. assert_type(c8.real, np.float32)
  21. assert_type(c8.imag, np.float32)
  22. assert_type(c8.real.real, np.float32)
  23. assert_type(c8.real.imag, np.float32)
  24. assert_type(c8.itemsize, int)
  25. assert_type(c8.shape, tuple[()])
  26. assert_type(c8.strides, tuple[()])
  27. assert_type(c8.ndim, Literal[0])
  28. assert_type(c8.size, Literal[1])
  29. assert_type(c8.squeeze(), np.complex64)
  30. assert_type(c8.byteswap(), np.complex64)
  31. assert_type(c8.transpose(), np.complex64)
  32. assert_type(c8.dtype, np.dtype[np.complex64])
  33. assert_type(c8.real, np.float32)
  34. assert_type(c16.imag, np.float64)
  35. assert_type(np.str_('foo'), np.str_)
  36. assert_type(V[0], Any)
  37. assert_type(V["field1"], Any)
  38. assert_type(V[["field1", "field2"]], np.void)
  39. V[0] = 5
  40. # Aliases
  41. assert_type(np.bool_(), np.bool[Literal[False]])
  42. assert_type(np.byte(), np.byte)
  43. assert_type(np.short(), np.short)
  44. assert_type(np.intc(), np.intc)
  45. assert_type(np.intp(), np.intp)
  46. assert_type(np.int_(), np.int_)
  47. assert_type(np.long(), np.long)
  48. assert_type(np.longlong(), np.longlong)
  49. assert_type(np.ubyte(), np.ubyte)
  50. assert_type(np.ushort(), np.ushort)
  51. assert_type(np.uintc(), np.uintc)
  52. assert_type(np.uintp(), np.uintp)
  53. assert_type(np.uint(), np.uint)
  54. assert_type(np.ulong(), np.ulong)
  55. assert_type(np.ulonglong(), np.ulonglong)
  56. assert_type(np.half(), np.half)
  57. assert_type(np.single(), np.single)
  58. assert_type(np.double(), np.double)
  59. assert_type(np.longdouble(), np.longdouble)
  60. assert_type(np.csingle(), np.csingle)
  61. assert_type(np.cdouble(), np.cdouble)
  62. assert_type(np.clongdouble(), np.clongdouble)
  63. assert_type(b.item(), bool)
  64. assert_type(i8.item(), int)
  65. assert_type(u8.item(), int)
  66. assert_type(f8.item(), float)
  67. assert_type(c16.item(), complex)
  68. assert_type(U.item(), str)
  69. assert_type(S.item(), bytes)
  70. assert_type(b.tolist(), bool)
  71. assert_type(i8.tolist(), int)
  72. assert_type(u8.tolist(), int)
  73. assert_type(f8.tolist(), float)
  74. assert_type(c16.tolist(), complex)
  75. assert_type(U.tolist(), str)
  76. assert_type(S.tolist(), bytes)
  77. assert_type(b.ravel(), np.ndarray[tuple[int], np.dtype[np.bool]])
  78. assert_type(i8.ravel(), np.ndarray[tuple[int], np.dtype[np.int64]])
  79. assert_type(u8.ravel(), np.ndarray[tuple[int], np.dtype[np.uint64]])
  80. assert_type(f8.ravel(), np.ndarray[tuple[int], np.dtype[np.float64]])
  81. assert_type(c16.ravel(), np.ndarray[tuple[int], np.dtype[np.complex128]])
  82. assert_type(U.ravel(), np.ndarray[tuple[int], np.dtype[np.str_]])
  83. assert_type(S.ravel(), np.ndarray[tuple[int], np.dtype[np.bytes_]])
  84. assert_type(b.flatten(), np.ndarray[tuple[int], np.dtype[np.bool]])
  85. assert_type(i8.flatten(), np.ndarray[tuple[int], np.dtype[np.int64]])
  86. assert_type(u8.flatten(), np.ndarray[tuple[int], np.dtype[np.uint64]])
  87. assert_type(f8.flatten(), np.ndarray[tuple[int], np.dtype[np.float64]])
  88. assert_type(c16.flatten(), np.ndarray[tuple[int], np.dtype[np.complex128]])
  89. assert_type(U.flatten(), np.ndarray[tuple[int], np.dtype[np.str_]])
  90. assert_type(S.flatten(), np.ndarray[tuple[int], np.dtype[np.bytes_]])
  91. assert_type(b.reshape(()), np.bool)
  92. assert_type(i8.reshape([]), np.int64)
  93. assert_type(b.reshape(1), np.ndarray[tuple[_1], np.dtype[np.bool]])
  94. assert_type(i8.reshape(-1), np.ndarray[tuple[_1], np.dtype[np.int64]])
  95. assert_type(u8.reshape(1, 1), np.ndarray[tuple[_1, _1], np.dtype[np.uint64]])
  96. assert_type(f8.reshape(1, -1), np.ndarray[tuple[_1, _1], np.dtype[np.float64]])
  97. assert_type(c16.reshape(1, 1, 1), np.ndarray[tuple[_1, _1, _1], np.dtype[np.complex128]])
  98. assert_type(U.reshape(1, 1, 1, 1), np.ndarray[tuple[_1, _1, _1, _1], np.dtype[np.str_]])
  99. assert_type(
  100. S.reshape(1, 1, 1, 1, 1),
  101. np.ndarray[
  102. # len(shape) >= 5
  103. tuple[_1, _1, _1, _1, _1, Unpack[tuple[_1, ...]]],
  104. np.dtype[np.bytes_],
  105. ],
  106. )
  107. assert_type(i8.astype(float), Any)
  108. assert_type(i8.astype(np.float64), np.float64)
  109. assert_type(i8.view(), np.int64)
  110. assert_type(i8.view(np.float64), np.float64)
  111. assert_type(i8.view(float), Any)
  112. assert_type(i8.view(np.float64, np.ndarray), np.float64)
  113. assert_type(i8.getfield(float), Any)
  114. assert_type(i8.getfield(np.float64), np.float64)
  115. assert_type(i8.getfield(np.float64, 8), np.float64)
  116. assert_type(f8.as_integer_ratio(), tuple[int, int])
  117. assert_type(f8.is_integer(), bool)
  118. assert_type(f8.__trunc__(), int)
  119. assert_type(f8.__getformat__("float"), str)
  120. assert_type(f8.hex(), str)
  121. assert_type(np.float64.fromhex("0x0.0p+0"), np.float64)
  122. assert_type(f8.__getnewargs__(), tuple[float])
  123. assert_type(c16.__getnewargs__(), tuple[float, float])
  124. assert_type(i8.numerator, np.int64)
  125. assert_type(i8.denominator, Literal[1])
  126. assert_type(u8.numerator, np.uint64)
  127. assert_type(u8.denominator, Literal[1])
  128. assert_type(m.numerator, np.timedelta64)
  129. assert_type(m.denominator, Literal[1])
  130. assert_type(round(i8), int)
  131. assert_type(round(i8, 3), np.int64)
  132. assert_type(round(u8), int)
  133. assert_type(round(u8, 3), np.uint64)
  134. assert_type(round(f8), int)
  135. assert_type(round(f8, 3), np.float64)
  136. assert_type(f8.__ceil__(), int)
  137. assert_type(f8.__floor__(), int)
  138. assert_type(i8.is_integer(), Literal[True])
  139. assert_type(O.real, np.object_)
  140. assert_type(O.imag, np.object_)
  141. assert_type(int(O), int)
  142. assert_type(float(O), float)
  143. assert_type(complex(O), complex)
  144. # These fail fail because of a mypy __new__ bug:
  145. # https://github.com/python/mypy/issues/15182
  146. # According to the typing spec, the following statements are valid, see
  147. # https://typing.readthedocs.io/en/latest/spec/constructors.html#new-method
  148. # assert_type(np.object_(), None)
  149. # assert_type(np.object_(None), None)
  150. # assert_type(np.object_(array_nd), np.ndarray[Any, np.dtype[np.object_]])
  151. # assert_type(np.object_([]), npt.NDArray[np.object_])
  152. # assert_type(np.object_(()), npt.NDArray[np.object_])
  153. # assert_type(np.object_(range(4)), npt.NDArray[np.object_])
  154. # assert_type(np.object_(+42), int)
  155. # assert_type(np.object_(1 / 137), float)
  156. # assert_type(np.object_('Developers! ' * (1 << 6)), str)
  157. # assert_type(np.object_(object()), object)
  158. # assert_type(np.object_({False, True, NotADirectoryError}), set[Any])
  159. # assert_type(np.object_({'spam': 'food', 'ham': 'food'}), dict[str, str])