rec.pyi 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. import io
  2. from typing import Any
  3. import numpy as np
  4. import numpy.typing as npt
  5. from typing_extensions import assert_type
  6. AR_i8: npt.NDArray[np.int64]
  7. REC_AR_V: np.recarray[tuple[int, ...], np.dtype[np.record]]
  8. AR_LIST: list[npt.NDArray[np.int64]]
  9. record: np.record
  10. file_obj: io.BufferedIOBase
  11. assert_type(np.rec.format_parser(
  12. formats=[np.float64, np.int64, np.bool],
  13. names=["f8", "i8", "?"],
  14. titles=None,
  15. aligned=True,
  16. ), np.rec.format_parser)
  17. assert_type(np.rec.format_parser.dtype, np.dtype[np.void])
  18. assert_type(record.field_a, Any)
  19. assert_type(record.field_b, Any)
  20. assert_type(record["field_a"], Any)
  21. assert_type(record["field_b"], Any)
  22. assert_type(record.pprint(), str)
  23. record.field_c = 5
  24. assert_type(REC_AR_V.field(0), Any)
  25. assert_type(REC_AR_V.field("field_a"), Any)
  26. assert_type(REC_AR_V.field(0, AR_i8), None)
  27. assert_type(REC_AR_V.field("field_a", AR_i8), None)
  28. assert_type(REC_AR_V["field_a"], npt.NDArray[Any])
  29. assert_type(REC_AR_V.field_a, Any)
  30. assert_type(REC_AR_V.__array_finalize__(object()), None)
  31. assert_type(
  32. np.recarray(
  33. shape=(10, 5),
  34. formats=[np.float64, np.int64, np.bool],
  35. order="K",
  36. byteorder="|",
  37. ),
  38. np.recarray[Any, np.dtype[np.record]],
  39. )
  40. assert_type(
  41. np.recarray(
  42. shape=(10, 5),
  43. dtype=[("f8", np.float64), ("i8", np.int64)],
  44. strides=(5, 5),
  45. ),
  46. np.recarray[Any, np.dtype[Any]],
  47. )
  48. assert_type(np.rec.fromarrays(AR_LIST), np.recarray[Any, np.dtype[Any]])
  49. assert_type(
  50. np.rec.fromarrays(AR_LIST, dtype=np.int64),
  51. np.recarray[Any, np.dtype[Any]],
  52. )
  53. assert_type(
  54. np.rec.fromarrays(
  55. AR_LIST,
  56. formats=[np.int64, np.float64],
  57. names=["i8", "f8"]
  58. ),
  59. np.recarray[Any, np.dtype[np.record]],
  60. )
  61. assert_type(
  62. np.rec.fromrecords((1, 1.5)),
  63. np.recarray[Any, np.dtype[np.record]]
  64. )
  65. assert_type(
  66. np.rec.fromrecords(
  67. [(1, 1.5)],
  68. dtype=[("i8", np.int64), ("f8", np.float64)],
  69. ),
  70. np.recarray[Any, np.dtype[np.record]],
  71. )
  72. assert_type(
  73. np.rec.fromrecords(
  74. REC_AR_V,
  75. formats=[np.int64, np.float64],
  76. names=["i8", "f8"]
  77. ),
  78. np.recarray[Any, np.dtype[np.record]],
  79. )
  80. assert_type(
  81. np.rec.fromstring(
  82. b"(1, 1.5)",
  83. dtype=[("i8", np.int64), ("f8", np.float64)],
  84. ),
  85. np.recarray[Any, np.dtype[np.record]],
  86. )
  87. assert_type(
  88. np.rec.fromstring(
  89. REC_AR_V,
  90. formats=[np.int64, np.float64],
  91. names=["i8", "f8"]
  92. ),
  93. np.recarray[Any, np.dtype[np.record]],
  94. )
  95. assert_type(np.rec.fromfile(
  96. "test_file.txt",
  97. dtype=[("i8", np.int64), ("f8", np.float64)],
  98. ), np.recarray[Any, np.dtype[Any]])
  99. assert_type(
  100. np.rec.fromfile(
  101. file_obj,
  102. formats=[np.int64, np.float64],
  103. names=["i8", "f8"]
  104. ),
  105. np.recarray[Any, np.dtype[np.record]],
  106. )
  107. assert_type(np.rec.array(AR_i8), np.recarray[Any, np.dtype[np.int64]])
  108. assert_type(
  109. np.rec.array([(1, 1.5)], dtype=[("i8", np.int64), ("f8", np.float64)]),
  110. np.recarray[Any, np.dtype[Any]],
  111. )
  112. assert_type(
  113. np.rec.array(
  114. [(1, 1.5)],
  115. formats=[np.int64, np.float64],
  116. names=["i8", "f8"]
  117. ),
  118. np.recarray[Any, np.dtype[np.record]],
  119. )
  120. assert_type(
  121. np.rec.array(
  122. None,
  123. dtype=np.float64,
  124. shape=(10, 3),
  125. ),
  126. np.recarray[Any, np.dtype[Any]],
  127. )
  128. assert_type(
  129. np.rec.array(
  130. None,
  131. formats=[np.int64, np.float64],
  132. names=["i8", "f8"],
  133. shape=(10, 3),
  134. ),
  135. np.recarray[Any, np.dtype[np.record]],
  136. )
  137. assert_type(
  138. np.rec.array(file_obj, dtype=np.float64),
  139. np.recarray[Any, np.dtype[Any]],
  140. )
  141. assert_type(
  142. np.rec.array(file_obj, formats=[np.int64, np.float64], names=["i8", "f8"]),
  143. np.recarray[Any, np.dtype[np.record]],
  144. )