8000 TYP: Type ``MaskedArray.all`` and ``MaskedArray.any`` (#28880) · numpy/numpy@cd0756a · GitHub
[go: up one dir, main page]
Skip to content

Commit cd0756a

Browse files
MarcoGorelliMarcoGorelli
authored
TYP: Type MaskedArray.all and MaskedArray.any (#28880)
1 parent b6cce1b commit cd0756a

File tree

3 files changed

+121
-2
lines changed

3 files changed

+121
-2
lines changed

‎numpy/ma/core.pyi

Lines changed: 91 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -487,8 +487,97 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
487487
def put(self, indices: _ArrayLikeInt_co, values: ArrayLike, mode: _ModeKind = "raise") -> None: ...
488488
def ids(self) -> tuple[int, int]: ...
489489
def iscontiguous(self) -> bool: ...
490-
def all(self, axis=..., out=..., keepdims=...): ...
491-
def any(self, axis=..., out=..., keepdims=...): ...
490+
491+
@overload
492+
def all(
493+
self,
494+
axis: None = None,
495+
out: None = None,
496+
keepdims: Literal[False] | _NoValueType = ...,
497+
) -> bool_: ...
498+
@overload
499+
def all(
500+
self,
501+
axis: _ShapeLike | None = None,
502+
out: None = None,
503+
*,
504+
keepdims: Literal[True],
505+
) -> _MaskedArray[bool_]: ...
506+
@overload
507+
def all(
508+
self,
509+
axis: _ShapeLike | None,
510+
out: None,
511+
keepdims: Literal[True],
512+
) -> _MaskedArray[bool_]: ...
513+
@overload
514+
def all(
515+
self,
516+
axis: _ShapeLike | None = None,
517+
out: None = None,
518+
keepdims: bool | _NoValueType = ...,
519+
) -> bool_ | _MaskedArray[bool_]: ...
520+
@overload
521+
def all(
522+
self,
523+
axis: _ShapeLike | None = None,
524+
*,
525+
out: _ArrayT,
526+
keepdims: bool | _NoValueType = ...,
527+
) -> _ArrayT: ...
528+
@overload
529+
def all(
530+
self,
531+
axis: _ShapeLike | None,
532+
out: _ArrayT,
533+
keepdims: bool | _NoValueType = ...,
534+
) -> _ArrayT: ...
535+
536+
@overload
537+
def any(
538+
self,
539+
axis: None = None,
540+
out: None = None,
541+
keepdims: Literal[False] | _NoValueType = ...,
542+
) -> bool_: ...
543+
@overload
544+
def any(
545+
self,
546+
axis: _ShapeLike | None = None,
547+
out: None = None,
548+
*,
549+
keepdims: Literal[True],
550+
) -> _MaskedArray[bool_]: ...
551+
@overload
552+
def any(
553+
self,
554+
axis: _ShapeLike | None,
555+
out: None,
556+
keepdims: Literal[True],
557+
) -> _MaskedArray[bool_]: ...
558+
@overload
559+
def any(
560+
self,
561+
axis: _ShapeLike | None = None,
562+
out: None = None,
563+
keepdims: bool | _NoValueType = ...,
564+
) -> bool_ | _MaskedArray[bool_]: ...
565+
@overload
566+
def any(
567+
self,
568+
axis: _ShapeLike | None = None,
569+
*,
570+
out: _ArrayT,
571+
keepdims: bool | _NoValueType = ...,
572+
) -> _ArrayT: ...
573+
@overload
574+
def any(
575+
self,
576+
axis: _ShapeLike | None,
577+
out: _ArrayT,
578+
keepdims: bool | _NoValueType = ...,
579+
) -> _ArrayT: ...
580+
492581
def nonzero(self) -> tuple[_Array1D[intp], *tuple[_Array1D[intp], ...]]: ...
493582
def trace(self, offset=..., axis1=..., axis2=..., dtype=..., out=...): ...
494583
def dot(self, b, out=..., strict=...): ...

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,14 @@ np.ma.argmax(m, keepdims=1.0) # E: No overload variant
6363
np.ma.argmax(m, out=1.0) # E: No overload variant
6464
np.ma.argmax(m, fill_value=lambda x: 27) # E: No overload variant
6565

66+
m.all(axis=1.0) # E: No overload variant
67+
m.all(keepdims=1.0) # E: No overload variant
68+
m.all(out=1.0) # E: No overload variant
69+
70+
m.any(axis=1.0) # E: No overload variant
71+
m.any(keepdims=1.0) # E: No overload variant
72+
m.any(out=1.0) # E: No overload variant
73+
6674
m.sort(axis=(0,1)) # E: No overload variant
6775
m.sort(axis=None) # E: No overload variant
6876
m.sort(kind='cabbage') # E: No overload variant

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,28 @@ assert_type(np.ma.argmax(MAR_b, keepdims=True), Any)
141141
assert_type(np.ma.argmax(MAR_f4, out=MAR_subclass), MaskedArraySubclass)
142142
assert_type(np.ma.argmax(MAR_f4, None, None, out=MAR_subclass), MaskedArraySubclass)
143143

144+
assert_type(MAR_b.all(), np.bool)
145+
assert_type(MAR_f4.all(), np.bool)
146+
assert_type(MAR_f4.all(keepdims=False), np.bool)
147+
assert_type(MAR_b.all(axis=0), np.bool | MaskedArray[np.bool])
148+
assert_type(MAR_b.all(axis=0, keepdims=True), MaskedArray[np.bool])
149+
assert_type(MAR_b.all(0, None, True), MaskedArray[np.bool])
150+
assert_type(MAR_f4.all(axis=0), np.bool | MaskedArray[np.bool])
151+
assert_type(MAR_b.all(keepdims=True), MaskedArray[np.bool])
152+
assert_type(MAR_f4.all(out=MAR_subclass), MaskedArraySubclass)
153+
assert_type(MAR_f4.all(None, out=MAR_subclass), MaskedArraySubclass)
154+
155+
assert_type(MAR_b.any(), np.bool)
156+
assert_type(MAR_f4.any(), np.bool)
157+
assert_type(MAR_f4.any(keepdims=False), np.bool)
158+
assert_type(MAR_b.any(axis=0), np.bool | MaskedArray[np.bool])
159+
assert_type(MAR_b.any(axis=0, keepdims=True), MaskedArray[np.bool])
160+
assert_type(MAR_b.any(0, None, True), MaskedArray[np.bool])
161+
assert_type(MAR_f4.any(axis=0), np.bool | MaskedArray[np.bool])
162+
assert_type(MAR_b.any(keepdims=True), MaskedArray[np.bool])
163+
assert_type(MAR_f4.any(out=MAR_subclass), MaskedArraySubclass)
164+
assert_type(MAR_f4.any(None, out=MAR_subclass), MaskedArraySubclass)
165+
144166
assert_type(MAR_f4.sort(), None)
145167
assert_type(MAR_f4.sort(axis=0, kind='quicksort', order='K', endwith=False, fill_value=42., stable=False), None)
146168

0 commit comments

Comments
 (0)
0