Repository: numpy/numpy
File: numpy/typing/tests/data/reveal/ma.pyi (Python)

Code before:
assert_type(MAR_dt64 <= AR_dt64, MaskedNDArray[np.bool])
assert_type(MAR_o <= AR_o, MaskedNDArray[np.bool])
assert_type(MAR_1d <= 0, MaskedNDArray[np.bool])
assert_type(MAR_s <= MAR_s, MaskedNDArray[np.bool])
assert_type(MAR_byte <= MAR_byte, MaskedNDArray[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[Any]])

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(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[Any]])

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[Any]])  # type: ignore[assert-type]


Review comment (suggestion):
```suggestion
assert_type(np.ma.compressed([[1,2,3]]), np.ndarray[tuple[int], np.dtype])
```

Code after:
assert_type(MAR_dt64 <= AR_dt64, MaskedNDArray[np.bool])
assert_type(MAR_o <= AR_o, MaskedNDArray[np.bool])
assert_type(MAR_1d <= 0, MaskedNDArray[np.bool])
assert_type(MAR_s <= MAR_s, MaskedNDArray[np.bool])
assert_type(MAR_byte <= MAR_byte, MaskedNDArray[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(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]