Skip to content

Commit 06c04a3

Browse files
MarcoGorelliMarcoGorelli
authored
TYP: Type MaskedArray.__{iadd,isub,imul,itruediv,ifloordiv,ipow}__ (numpy#28986)
1 parent 0506cf6 commit 06c04a3

File tree

3 files changed

+353
-12
lines changed

3 files changed

+353
-12
lines changed

‎numpy/ma/core.pyi‎

Lines changed: 173 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,24 @@ from numpy import (
1717
amax,
1818
amin,
1919
bool_,
20+
bytes_,
21+
character,
22+
complexfloating,
23+
datetime64,
2024
dtype,
25+
dtypes,
2126
expand_dims,
2227
float64,
28+
floating,
2329
generic,
2430
int_,
2531
intp,
2632
ndarray,
33+
object_,
34+
signedinteger,
35+
str_,
36+
timedelta64,
37+
unsignedinteger,
2738
)
2839
from numpy._globals import _NoValueType
2940
from numpy._typing import (
@@ -32,8 +43,15 @@ from numpy._typing import (
3243
_AnyShape,
3344
_ArrayLike,
3445
_ArrayLikeBool_co,
46+
_ArrayLikeBytes_co,
47+
_ArrayLikeComplex_co,
48+
_ArrayLikeFloat_co,
3549
_ArrayLikeInt,
3650
_ArrayLikeInt_co,
51+
_ArrayLikeStr_co,
52+
_ArrayLikeString_co,
53+
_ArrayLikeTD64_co,
54+
_ArrayLikeUInt_co,
3755
_DTypeLikeBool,
3856
_IntLike_co,
3957
_ScalarLike_co,
@@ -457,12 +475,161 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
457475
def __rfloordiv__(self, other): ...
458476
def __pow__(self, other, mod: None = None, /): ...
459477
def __rpow__(self, other, mod: None = None, /): ...
460-
def __iadd__(self, other): ...
461-
def __isub__(self, other): ...
462-
def __imul__(self, other): ...
463-
def __ifloordiv__(self, other): ...
464-
def __itruediv__(self, other): ...
465-
def __ipow__(self, other): ...
478+
479+
# Keep in sync with `ndarray.__iadd__`, except that `_MaskedArray[unsignedinteger]` does not accept
480+
# _IntLike_co for `other`.
481+
@overload
482+
def __iadd__(
483+
self: _MaskedArray[np.bool], other: _ArrayLikeBool_co, /
484+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
485+
@overload
486+
def __iadd__(
487+
self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
488+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
489+
@overload
490+
def __iadd__(
491+
self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
492+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
493+
@overload
494+
def __iadd__(
495+
self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
496+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
497+
@overload
498+
def __iadd__(
499+
self: _MaskedArray[complexfloating], other: _ArrayLikeComplex_co, /
500+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
501+
@overload
502+
def __iadd__(
503+
self: _MaskedArray[timedelta64 | datetime64], other: _ArrayLikeTD64_co, /
504+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
505+
@overload
506+
def __iadd__(self: _MaskedArray[bytes_], other: _ArrayLikeBytes_co, /) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
507+
@overload
508+
def __iadd__(
509+
self: MaskedArray[Any, dtype[str_] | dtypes.StringDType],
510+
other: _ArrayLikeStr_co | _ArrayLikeString_co,
511+
/,
512+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
513+
@overload
514+
def __iadd__(
515+
self: _MaskedArray[object_], other: Any, /
516+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
517+
518+
# Keep in sync with `ndarray.__isub__`, except that `_MaskedArray[unsignedinteger]` does not accept
519+
# _IntLike_co for `other`.
520+
@overload
521+
def __isub__(
522+
self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
523+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
524+
@overload
525+
def __isub__(
526+
self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
527+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
528+
@overload
529+
def __isub__(
530+
self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
531+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
532+
@overload
533+
def __isub__(
534+
self: _MaskedArray[complexfloating], other: _ArrayLikeComplex_co, /
535+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
536+
@overload
537+
def __isub__(
538+
self: _MaskedArray[timedelta64 | datetime64], other: _ArrayLikeTD64_co, /
539+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
540+
@overload
541+
def __isub__(
542+
self: _MaskedArray[object_], other: Any, /
543+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
544+
545+
# Keep in sync with `ndarray.__imul__`, except that `_MaskedArray[unsignedinteger]` does not accept
546+
# _IntLike_co for `other`.
547+
@overload
548+
def __imul__(
549+
self: _MaskedArray[np.bool], other: _ArrayLikeBool_co, /
550+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
551+
@overload
552+
def __imul__(
553+
self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
554+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
555+
@overload
556+
def __imul__(
557+
self: MaskedArray[Any, dtype[signedinteger] | dtype[character] | dtypes.StringDType],
558+
other: _ArrayLikeInt_co, /
559+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
560+
@overload
561+
def __imul__(
562+
self: _MaskedArray[floating | timedelta64], other: _ArrayLikeFloat_co, /
563+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
564+
@overload
565+
def __imul__(
566+
self: _MaskedArray[complexfloating], other: _ArrayLikeComplex_co, /
567+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
568+
@overload
569+
def __imul__(
570+
self: _MaskedArray[object_], other: Any, /
571+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
572+
573+
# Keep in sync with `ndarray.__ifloordiv__`, except that `_MaskedArray[unsignedinteger]` does not accept
574+
# _IntLike_co for `other`.
575+
@overload
576+
def __ifloordiv__(
577+
self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
578+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
579+
@overload
580+
def __ifloordiv__(
581+
self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
582+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
583+
@overload
584+
def __ifloordiv__(
585+
self: _MaskedArray[floating | timedelta64], other: _ArrayLikeFloat_co, /
586+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
587+
@overload
588+
def __ifloordiv__(
589+
self: _MaskedArray[object_], other: Any, /
590+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
591+
592+
# Keep in sync with `ndarray.__itruediv__`, except that `_MaskedArray[unsignedinteger]` does not accept
593+
# _IntLike_co for `other`.
594+
@overload
595+
def __itruediv__(
596+
self: _MaskedArray[floating | timedelta64], other: _ArrayLikeFloat_co, /
597+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
598+
@overload
599+
def __itruediv__(
600+
self: _MaskedArray[complexfloating],
601+
other: _ArrayLikeComplex_co,
602+
/,
603+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
604+
@overload
605+
def __itruediv__(
606+
self: _MaskedArray[object_], other: Any, /
607+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
608+
609+
# Keep in sync with `ndarray.__ipow__`, except that `_MaskedArray[unsignedinteger]` does not accept
610+
# _IntLike_co for `other`.
611+
@overload
612+
def __ipow__(
613+
self: _MaskedArray[unsignedinteger], other: _ArrayLikeUInt_co, /
614+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
615+
@overload
616+
def __ipow__(
617+
self: _MaskedArray[signedinteger], other: _ArrayLikeInt_co, /
618+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
619+
@overload
620+
def __ipow__(
621+
self: _MaskedArray[floating], other: _ArrayLikeFloat_co, /
622+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
623+
@overload
624+
def __ipow__(
625+
self: _MaskedArray[complexfloating], other: _ArrayLikeComplex_co, /
626+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
627+
@overload
628+
def __ipow__(
629+
self: _MaskedArray[object_], other: Any, /
630+
) -> MaskedArray[_ShapeT_co, _DTypeT_co]: ...
631+
632+
#
466633
@property # type: ignore[misc]
467634
def imag(self: _HasDTypeWithRealAndImag[object, _ScalarT], /) -> MaskedArray[_ShapeT_co, dtype[_ScalarT]]: ...
468635
get_imag: Any

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,16 @@
1+
from typing import TypeAlias, TypeVar
2+
13
import numpy as np
24
import numpy.typing as npt
5+
from numpy._typing import _Shape
6+
7+
_ScalarT = TypeVar("_ScalarT", bound=np.generic)
8+
MaskedArray: TypeAlias = np.ma.MaskedArray[_Shape, np.dtype[_ScalarT]]
39

410
MAR_1d_f8: np.ma.MaskedArray[tuple[int], np.dtype[np.float64]]
11+
MAR_b: MaskedArray[np.bool]
12+
MAR_c: MaskedArray[np.complex128]
13+
MAR_td64: MaskedArray[np.timedelta64]
514

615
AR_b: npt.NDArray[np.bool]
716

@@ -127,4 +136,8 @@ np.ma.allclose(MAR_1d_f8, [1,2,3], atol='.5') # type: ignore[arg-type]
127136

128137
MAR_1d_f8.__setmask__('mask') # type: ignore[arg-type]
129138

139+
MAR_b *= 2 # type: ignore[arg-type]
140+
MAR_c //= 2 # type: ignore[misc]
141+
MAR_td64 **= 2 # type: ignore[misc]
142+
130143
MAR_1d_f8.swapaxes(axis1=1, axis2=0) # type: ignore[call-arg]

0 commit comments

Comments
 (0)