codekingpro/portable-devtools

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