8000 TYP: Type ``MaskedArray.compressed`` and ``np.ma.compressed`` · numpy/numpy@525c79b · GitHub
[go: up one dir, main page]
Skip to content

Commit 525c79b

Browse files
MarcoGorelliMarcoGorelli
committed
TYP: Type MaskedArray.compressed and np.ma.compressed 
1 parent 0c33694 commit 525c79b

File tree

3 files changed

+12
-2
lines changed

3 files changed

+12
-2
lines changed

‎numpy/ma/core.pyi

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -427,7 +427,7 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
427427
get_fill_value: Any
428428
set_fill_value: Any
429429
def filled(self, /, fill_value: _ScalarLike_co | None = None) -> ndarray[_ShapeT_co, _DTypeT_co]: ...
430-
def compressed(self): ...
430+
def compressed(self: MaskedArray[Any, _DTypeT_co]) -> ndarray[tuple[int], _DTypeT_co]: ...
431431
def compress(self, condition, axis=..., out=...): ...
432432
def __eq__(self, other): ...
433433
def __ne__(self, other): ...
@@ -1154,7 +1154,10 @@ def sort(
11541154
*,
11551155
stable: Literal[False] | None = False,
11561156
) -> NDArray[Any]: ...
1157-
def compressed(x): ...
1157+
@overload
1158+
def compressed(x: _ArrayLike[_ScalarT_co]) -> ndarray[tuple[int], dtype[_ScalarT_co]]: ...
1159+
@overload
1160+
def compressed(x: ArrayLike) -> ndarray[tuple[int], Any]: ...
11581161
def concatenate(arrays, axis=...): ...
11591162
def diag(v, k=...): ...
11601163
def left_shift(a, n): ...

‎numpy/typing/tests/data/fail/ma.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,3 +114,5 @@ m.put(4, 999, mode='flip') # E: No overload variant
114114
np.ma.put(m, 4, 999, mode='flip') # E: No overload variant
115115

116116
np.ma.put([1,1,3], 0, 999) # E: No overload variant
117+
118+
np.ma.compressed(lambda: 'compress me') # E: No overload variant

‎numpy/typing/tests/data/reveal/ma.pyi

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,11 @@ assert_type(np.ma.count(MAR_o, keepdims=True), NDArray[np.int_])
249249
assert_type(np.ma.count(MAR_o, axis=None, keepdims=True), NDArray[np.int_])
250250
assert_type(np.ma.count(MAR_o, None, True), NDArray[np.int_])
251251

252+
assert_type(MAR_f4.compressed(), np.ndarray[tuple[int], np.dtype[np.float32]])
253+
254+
assert_type(np.ma.compressed(MAR_i8), np.ndarray[tuple[int], np.dtype[np.int64]])
255+
assert_type(np.ma.compressed([[1,2,3]]), np.ndarray[tuple[int], np.dtype[Any]])
256+
252257
assert_type(MAR_f4.put([0,4,8], [10,20,30]), None)
253258
assert_type(MAR_f4.put(4, 999), None)
254259
assert_type(MAR_f4.put(4, 999, mode='clip'), None)

0 commit comments

Comments
 (0)
0