1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
|
from typing import Any, Literal, TypeAlias, TypeVar, assert_type
import numpy as np
from numpy import dtype, generic
from numpy._typing import NDArray, _AnyShape
_ScalarT = TypeVar("_ScalarT", bound=generic)
MaskedArray: TypeAlias = np.ma.MaskedArray[_AnyShape, dtype[_ScalarT]]
_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_ScalarT]]
class MaskedArraySubclass(MaskedArray[np.complex128]): ...
AR_b: NDArray[np.bool]
AR_f4: NDArray[np.float32]
AR_dt64: NDArray[np.datetime64]
AR_td64: NDArray[np.timedelta64]
AR_o: NDArray[np.timedelta64]
MAR_c16: MaskedArray[np.complex128]
MAR_b: MaskedArray[np.bool]
MAR_f4: MaskedArray[np.float32]
MAR_f8: MaskedArray[np.float64]
MAR_i8: MaskedArray[np.int64]
MAR_dt64: MaskedArray[np.datetime64]
MAR_td64: MaskedArray[np.timedelta64]
MAR_o: MaskedArray[np.object_]
MAR_s: MaskedArray[np.str_]
MAR_byte: MaskedArray[np.bytes_]
MAR_V: MaskedArray[np.void]
MAR_subclass: MaskedArraySubclass
MAR_1d: np.ma.MaskedArray[tuple[int], np.dtype]
MAR_2d_f4: np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]]
b: np.bool
f4: np.float32
f: float
assert_type(MAR_1d.shape, tuple[int])
assert_type(MAR_f4.dtype, np.dtype[np.float32])
assert_type(int(MAR_i8), int)
assert_type(float(MAR_f4), float)
assert_type(np.ma.min(MAR_b), np.bool)
assert_type(np.ma.min(MAR_f4), np.float32)
assert_type(np.ma.min(MAR_b, axis=0), Any)
assert_type(np.ma.min(MAR_f4, axis=0), Any)
assert_type(np.ma.min(MAR_b, keepdims=True), Any)
assert_type(np.ma.min(MAR_f4, keepdims=True), Any)
assert_type(np.ma.min(MAR_f4, out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.min(MAR_f4, 0, MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.min(MAR_f4, None, MAR_subclass), MaskedArraySubclass)
assert_type(MAR_b.min(), np.bool)
assert_type(MAR_f4.min(), np.float32)
assert_type(MAR_b.min(axis=0), Any)
assert_type(MAR_f4.min(axis=0), Any)
assert_type(MAR_b.min(keepdims=True), Any)
assert_type(MAR_f4.min(keepdims=True), Any)
assert_type(MAR_f4.min(out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.min(0, MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.min(None, MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.max(MAR_b), np.bool)
assert_type(np.ma.max(MAR_f4), np.float32)
assert_type(np.ma.max(MAR_b, axis=0), Any)
assert_type(np.ma.max(MAR_f4, axis=0), Any)
assert_type(np.ma.max(MAR_b, keepdims=True), Any)
assert_type(np.ma.max(MAR_f4, keepdims=True), Any)
assert_type(np.ma.max(MAR_f4, out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.max(MAR_f4, 0, MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.max(MAR_f4, None, MAR_subclass), MaskedArraySubclass)
assert_type(MAR_b.max(), np.bool)
assert_type(MAR_f4.max(), np.float32)
assert_type(MAR_b.max(axis=0), Any)
assert_type(MAR_f4.max(axis=0), Any)
assert_type(MAR_b.max(keepdims=True), Any)
assert_type(MAR_f4.max(keepdims=True), Any)
assert_type(MAR_f4.max(out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.max(0, MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.max(None, MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.ptp(MAR_b), np.bool)
assert_type(np.ma.ptp(MAR_f4), np.float32)
assert_type(np.ma.ptp(MAR_b, axis=0), Any)
assert_type(np.ma.ptp(MAR_f4, axis=0), Any)
assert_type(np.ma.ptp(MAR_b, keepdims=True), Any)
assert_type(np.ma.ptp(MAR_f4, keepdims=True), Any)
assert_type(np.ma.ptp(MAR_f4, out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.ptp(MAR_f4, 0, MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.ptp(MAR_f4, None, MAR_subclass), MaskedArraySubclass)
assert_type(MAR_b.ptp(), np.bool)
assert_type(MAR_f4.ptp(), np.float32)
assert_type(MAR_b.ptp(axis=0), Any)
assert_type(MAR_f4.ptp(axis=0), Any)
assert_type(MAR_b.ptp(keepdims=True), Any)
assert_type(MAR_f4.ptp(keepdims=True), Any)
assert_type(MAR_f4.ptp(out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.ptp(0, MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.ptp(None, MAR_subclass), MaskedArraySubclass)
assert_type(MAR_b.argmin(), np.intp)
assert_type(MAR_f4.argmin(), np.intp)
assert_type(MAR_f4.argmax(fill_value=6.28318, keepdims=False), np.intp)
assert_type(MAR_b.argmin(axis=0), Any)
assert_type(MAR_f4.argmin(axis=0), Any)
assert_type(MAR_b.argmin(keepdims=True), Any)
assert_type(MAR_f4.argmin(out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.argmin(None, None, out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.argmin(MAR_b), np.intp)
assert_type(np.ma.argmin(MAR_f4), np.intp)
assert_type(np.ma.argmin(MAR_f4, fill_value=6.28318, keepdims=False), np.intp)
assert_type(np.ma.argmin(MAR_b, axis=0), Any)
assert_type(np.ma.argmin(MAR_f4, axis=0), Any)
assert_type(np.ma.argmin(MAR_b, keepdims=True), Any)
assert_type(np.ma.argmin(MAR_f4, out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.argmin(MAR_f4, None, None, out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_b.argmax(), np.intp)
assert_type(MAR_f4.argmax(), np.intp)
assert_type(MAR_f4.argmax(fill_value=6.28318, keepdims=False), np.intp)
assert_type(MAR_b.argmax(axis=0), Any)
assert_type(MAR_f4.argmax(axis=0), Any)
assert_type(MAR_b.argmax(keepdims=True), Any)
assert_type(MAR_f4.argmax(out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.argmax(None, None, out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.argmax(MAR_b), np.intp)
assert_type(np.ma.argmax(MAR_f4), np.intp)
assert_type(np.ma.argmax(MAR_f4, fill_value=6.28318, keepdims=False), np.intp)
assert_type(np.ma.argmax(MAR_b, axis=0), Any)
assert_type(np.ma.argmax(MAR_f4, axis=0), Any)
assert_type(np.ma.argmax(MAR_b, keepdims=True), Any)
assert_type(np.ma.argmax(MAR_f4, out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.argmax(MAR_f4, None, None, out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_b.all(), np.bool)
assert_type(MAR_f4.all(), np.bool)
assert_type(MAR_f4.all(keepdims=False), np.bool)
assert_type(MAR_b.all(axis=0), np.bool | MaskedArray[np.bool])
assert_type(MAR_b.all(axis=0, keepdims=True), MaskedArray[np.bool])
assert_type(MAR_b.all(0, None, True), MaskedArray[np.bool])
assert_type(MAR_f4.all(axis=0), np.bool | MaskedArray[np.bool])
assert_type(MAR_b.all(keepdims=True), MaskedArray[np.bool])
assert_type(MAR_f4.all(out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.all(None, out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_b.any(), np.bool)
assert_type(MAR_f4.any(), np.bool)
assert_type(MAR_f4.any(keepdims=False), np.bool)
assert_type(MAR_b.any(axis=0), np.bool | MaskedArray[np.bool])
assert_type(MAR_b.any(axis=0, keepdims=True), MaskedArray[np.bool])
assert_type(MAR_b.any(0, None, True), MaskedArray[np.bool])
assert_type(MAR_f4.any(axis=0), np.bool | MaskedArray[np.bool])
assert_type(MAR_b.any(keepdims=True), MaskedArray[np.bool])
assert_type(MAR_f4.any(out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.any(None, out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f4.sort(), None)
assert_type(MAR_f4.sort(axis=0, kind='quicksort', order='K', endwith=False, fill_value=42., stable=False), None)
assert_type(np.ma.sort(MAR_f4), MaskedArray[np.float32])
assert_type(np.ma.sort(MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.sort([[0, 1], [2, 3]]), NDArray[Any])
assert_type(np.ma.sort(AR_f4), NDArray[np.float32])
assert_type(MAR_f8.take(0), np.float64)
assert_type(MAR_1d.take(0), Any)
assert_type(MAR_f8.take([0]), MaskedArray[np.float64])
assert_type(MAR_f8.take(0, out=MAR_subclass), MaskedArraySubclass)
assert_type(MAR_f8.take([0], out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.take(f, 0), Any)
assert_type(np.ma.take(f4, 0), np.float32)
assert_type(np.ma.take(MAR_f8, 0), np.float64)
assert_type(np.ma.take(AR_f4, 0), np.float32)
assert_type(np.ma.take(MAR_1d, 0), Any)
assert_type(np.ma.take(MAR_f8, [0]), MaskedArray[np.float64])
assert_type(np.ma.take(AR_f4, [0]), MaskedArray[np.float32])
assert_type(np.ma.take(MAR_f8, 0, out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.take(MAR_f8, [0], out=MAR_subclass), MaskedArraySubclass)
assert_type(np.ma.take([1], [0]), MaskedArray[Any])
assert_type(np.ma.take(np.eye(2), 1, axis=0), MaskedArray[np.float64])
assert_type(MAR_f4.partition(1), None)
assert_type(MAR_V.partition(1, axis=0, kind='introselect', order='K'), None)
assert_type(MAR_f4.argpartition(1), MaskedArray[np.intp])
assert_type(MAR_1d.argpartition(1, axis=0, kind='introselect', order='K'), MaskedArray[np.intp])
assert_type(np.ma.ndim(f4), int)
assert_type(np.ma.ndim(MAR_b), int)
assert_type(np.ma.ndim(AR_f4), int)
assert_type(np.ma.size(b), int)
assert_type(np.ma.size(MAR_f4, axis=0), int)
assert_type(np.ma.size(AR_f4), int)
assert_type(np.ma.is_masked(MAR_f4), bool)
assert_type(MAR_f4.ids(), tuple[int, int])
assert_type(MAR_f4.iscontiguous(), bool)
assert_type(MAR_f4 >= 3, MaskedArray[np.bool])
assert_type(MAR_i8 >= AR_td64, MaskedArray[np.bool])
assert_type(MAR_b >= AR_td64, MaskedArray[np.bool])
assert_type(MAR_td64 >= AR_td64, MaskedArray[np.bool])
assert_type(MAR_dt64 >= AR_dt64, MaskedArray[np.bool])
assert_type(MAR_o >= AR_o, MaskedArray[np.bool])
assert_type(MAR_1d >= 0, MaskedArray[np.bool])
assert_type(MAR_s >= MAR_s, MaskedArray[np.bool])
assert_type(MAR_byte >= MAR_byte, MaskedArray[np.bool])
assert_type(MAR_f4 > 3, MaskedArray[np.bool])
assert_type(MAR_i8 > AR_td64, MaskedArray[np.bool])
assert_type(MAR_b > AR_td64, MaskedArray[np.bool])
assert_type(MAR_td64 > AR_td64, MaskedArray[np.bool])
assert_type(MAR_dt64 > AR_dt64, MaskedArray[np.bool])
assert_type(MAR_o > AR_o, MaskedArray[np.bool])
assert_type(MAR_1d > 0, MaskedArray[np.bool])
assert_type(MAR_s > MAR_s, MaskedArray[np.bool])
assert_type(MAR_byte > MAR_byte, MaskedArray[np.bool])
assert_type(MAR_f4 <= 3, MaskedArray[np.bool])
assert_type(MAR_i8 <= AR_td64, MaskedArray[np.bool])
assert_type(MAR_b <= AR_td64, MaskedArray[np.bool])
assert_type(MAR_td64 <= AR_td64, MaskedArray[np.bool])
assert_type(MAR_dt64 <= AR_dt64, MaskedArray[np.bool])
assert_type(MAR_o <= AR_o, MaskedArray[np.bool])
assert_type(MAR_1d <= 0, MaskedArray[np.bool])
assert_type(MAR_s <= MAR_s, MaskedArray[np.bool])
assert_type(MAR_byte <= MAR_byte, MaskedArray[np.bool])
assert_type(MAR_f4 < 3, MaskedArray[np.bool])
assert_type(MAR_i8 < AR_td64, MaskedArray[np.bool])
assert_type(MAR_b < AR_td64, MaskedArray[np.bool])
assert_type(MAR_td64 < AR_td64, MaskedArray[np.bool])
assert_type(MAR_dt64 < AR_dt64, MaskedArray[np.bool])
assert_type(MAR_o < AR_o, MaskedArray[np.bool])
assert_type(MAR_1d < 0, MaskedArray[np.bool])
assert_type(MAR_s < MAR_s, MaskedArray[np.bool])
assert_type(MAR_byte < MAR_byte, MaskedArray[np.bool])
assert_type(MAR_f4 <= 3, MaskedArray[np.bool])
assert_type(MAR_i8 <= AR_td64, MaskedArray[np.bool])
assert_type(MAR_b <= AR_td64, MaskedArray[np.bool])
assert_type(MAR_td64 <= AR_td64, MaskedArray[np.bool])
assert_type(MAR_dt64 <= AR_dt64, MaskedArray[np.bool])
assert_type(MAR_o <= AR_o, MaskedArray[np.bool])
assert_type(MAR_1d <= 0, MaskedArray[np.bool])
assert_type(MAR_s <= MAR_s, MaskedArray[np.bool])
assert_type(MAR_byte <= MAR_byte, MaskedArray[np.bool])
assert_type(MAR_byte.count(), int)
assert_type(MAR_f4.count(axis=None), int)
assert_type(MAR_f4.count(axis=0), NDArray[np.int_])
assert_type(MAR_b.count(axis=(0,1)), NDArray[np.int_])
assert_type(MAR_o.count(keepdims=True), NDArray[np.int_])
assert_type(MAR_o.count(axis=None, keepdims=True), NDArray[np.int_])
assert_type(MAR_o.count(None, True), NDArray[np.int_])
assert_type(np.ma.count(MAR_byte), int)
assert_type(np.ma.count(MAR_byte, axis=None), int)
assert_type(np.ma.count(MAR_f4, axis=0), NDArray[np.int_])
assert_type(np.ma.count(MAR_b, axis=(0,1)), NDArray[np.int_])
assert_type(np.ma.count(MAR_o, keepdims=True), NDArray[np.int_])
assert_type(np.ma.count(MAR_o, axis=None, keepdims=True), NDArray[np.int_])
assert_type(np.ma.count(MAR_o, None, True), NDArray[np.int_])
assert_type(MAR_f4.compressed(), np.ndarray[tuple[int], np.dtype[np.float32]])
assert_type(np.ma.compressed(MAR_i8), np.ndarray[tuple[int], np.dtype[np.int64]])
assert_type(np.ma.compressed([[1,2,3]]), np.ndarray[tuple[int], np.dtype])
assert_type(MAR_f4.put([0,4,8], [10,20,30]), None)
assert_type(MAR_f4.put(4, 999), None)
assert_type(MAR_f4.put(4, 999, mode='clip'), None)
assert_type(np.ma.put(MAR_f4, [0,4,8], [10,20,30]), None)
assert_type(np.ma.put(MAR_f4, 4, 999), None)
assert_type(np.ma.put(MAR_f4, 4, 999, mode='clip'), None)
assert_type(np.ma.putmask(MAR_f4, [True, False], [0, 1]), None)
assert_type(np.ma.putmask(MAR_f4, np.False_, [0, 1]), None)
assert_type(MAR_f4.filled(float('nan')), NDArray[np.float32])
assert_type(MAR_i8.filled(), NDArray[np.int64])
assert_type(MAR_1d.filled(), np.ndarray[tuple[int], np.dtype])
assert_type(np.ma.filled(MAR_f4, float('nan')), NDArray[np.float32])
assert_type(np.ma.filled([[1,2,3]]), NDArray[Any])
# PyRight detects this one correctly, but mypy doesn't.
# https://github.com/numpy/numpy/pull/28742#discussion_r2048968375
assert_type(np.ma.filled(MAR_1d), np.ndarray[tuple[int], np.dtype]) # type: ignore[assert-type]
assert_type(MAR_b.repeat(3), np.ma.MaskedArray[tuple[int], np.dtype[np.bool]])
assert_type(MAR_2d_f4.repeat(MAR_i8), np.ma.MaskedArray[tuple[int], np.dtype[np.float32]])
assert_type(MAR_2d_f4.repeat(MAR_i8, axis=None), np.ma.MaskedArray[tuple[int], np.dtype[np.float32]])
assert_type(MAR_2d_f4.repeat(MAR_i8, axis=0), MaskedArray[np.float32])
assert_type(np.ma.allequal(AR_f4, MAR_f4), bool)
assert_type(np.ma.allequal(AR_f4, MAR_f4, fill_value=False), bool)
assert_type(np.ma.allclose(AR_f4, MAR_f4), bool)
assert_type(np.ma.allclose(AR_f4, MAR_f4, masked_equal=False), bool)
assert_type(np.ma.allclose(AR_f4, MAR_f4, rtol=.4, atol=.3), bool)
assert_type(MAR_2d_f4.ravel(), np.ma.MaskedArray[tuple[int], np.dtype[np.float32]])
assert_type(MAR_1d.ravel(order='A'), np.ma.MaskedArray[tuple[int], np.dtype[Any]])
assert_type(np.ma.getmask(MAR_f4), NDArray[np.bool] | np.bool)
# PyRight detects this one correctly, but mypy doesn't:
# `Revealed type is "Union[numpy.ndarray[Any, Any], numpy.bool[Any]]"`
assert_type(np.ma.getmask(MAR_1d), np.ndarray[tuple[int], np.dtype[np.bool]] | np.bool) # type: ignore[assert-type]
assert_type(np.ma.getmask(MAR_2d_f4), np.ndarray[tuple[int, int], np.dtype[np.bool]] | np.bool)
assert_type(np.ma.getmask([1,2]), NDArray[np.bool] | np.bool)
assert_type(np.ma.getmask(np.int64(1)), np.bool)
assert_type(np.ma.is_mask(MAR_1d), bool)
assert_type(np.ma.is_mask(AR_b), bool)
def func(x: object) -> None:
if np.ma.is_mask(x):
assert_type(x, NDArray[np.bool])
else:
assert_type(x, object)
assert_type(MAR_2d_f4.mT, np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(MAR_c16.real, MaskedArray[np.float64])
assert_type(MAR_c16.imag, MaskedArray[np.float64])
assert_type(MAR_2d_f4.baseclass, type[NDArray[Any]])
assert_type(MAR_b.swapaxes(0, 1), MaskedArray[np.bool])
assert_type(MAR_2d_f4.swapaxes(1, 0), MaskedArray[np.float32])
assert_type(np.ma.nomask, np.bool[Literal[False]])
assert_type(np.ma.MaskType, type[np.bool])
assert_type(MAR_1d.__setmask__([True, False]), None)
assert_type(MAR_1d.__setmask__(np.False_), None)
assert_type(MAR_2d_f4.harden_mask(), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(MAR_i8.harden_mask(), MaskedArray[np.int64])
assert_type(MAR_2d_f4.soften_mask(), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(MAR_i8.soften_mask(), MaskedArray[np.int64])
assert_type(MAR_f4.unshare_mask(), MaskedArray[np.float32])
assert_type(MAR_b.shrink_mask(), MaskedArray[np.bool_])
assert_type(MAR_i8.hardmask, bool)
assert_type(MAR_i8.sharedmask, bool)
assert_type(MAR_b.transpose(), MaskedArray[np.bool])
assert_type(MAR_2d_f4.transpose(), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(MAR_2d_f4.transpose(1, 0), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(MAR_2d_f4.transpose((1, 0)), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(MAR_b.T, MaskedArray[np.bool])
assert_type(MAR_2d_f4.T, np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(MAR_2d_f4.nonzero(), tuple[_Array1D[np.intp], *tuple[_Array1D[np.intp], ...]])
assert_type(MAR_2d_f4.nonzero()[0], _Array1D[np.intp])
|
