Skip to content

TYP: Type MaskedArray.__{add,radd,sub,rsub}__ #29012

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

TYP: Type MaskedArray.__{add,radd,sub,rsub}__ #29012

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9bc7b74
TYP: Type ``MaskedArray.__{add,radd,sub,rsub}__``
MarcoGorelli May 20, 2025
02fbee2
Merge remote-tracking branch 'upstream/main' into ma-arithmetic
MarcoGorelli May 20, 2025
5889f5c
"in-sync" comment in test, copy comment in __radd__, use _AnyShape
MarcoGorelli May 20, 2025
d605c2f
noop
MarcoGorelli May 21, 2025
dd7293b
Merge remote-tracking branch 'upstream/main' into ma-arithmetic
MarcoGorelli May 22, 2025
93edb06
Merge remote-tracking branch 'upstream/main' into ma-arithmetic
MarcoGorelli May 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
TYP: Type MaskedArray.__{add,radd,sub,rsub}__
  • Loading branch information
@MarcoGorelli
MarcoGorelli committed May 20, 2025
commit 9bc7b74b0f790d4c7ea7e25705277f1de679a440
202 changes: 197 additions & 5 deletions numpy/ma/core.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
# ruff: noqa: ANN001, ANN002, ANN003, ANN201, ANN202 ANN204, ANN401

from collections.abc import Sequence
from typing import Any, Literal, Self, SupportsIndex, TypeAlias, overload
from typing import Any, Literal, NoReturn, Self, SupportsIndex, TypeAlias, overload

from _typeshed import Incomplete
from typing_extensions import TypeIs, TypeVar, deprecated
Expand All @@ -19,17 +19,23 @@ from numpy import (
bool_,
bytes_,
character,
complex128,
complexfloating,
datetime64,
dtype,
dtypes,
expand_dims,
float16,
float32,
float64,
floating,
generic,
inexact,
int_,
integer,
intp,
ndarray,
number,
object_,
signedinteger,
str_,
Expand All @@ -40,14 +46,21 @@ from numpy._globals import _NoValueType
from numpy._typing import (
ArrayLike,
NDArray,
_32Bit,
_64Bit,
_AnyShape,
_ArrayLike,
_ArrayLikeBool_co,
_ArrayLikeBytes_co,
_ArrayLikeComplex128_co,
_ArrayLikeComplex_co,
_ArrayLikeDT64_co,
_ArrayLikeFloat64_co,
_ArrayLikeFloat_co,
_ArrayLikeInt,
_ArrayLikeInt_co,
_ArrayLikeNumber_co,
_ArrayLikeObject_co,
_ArrayLikeStr_co,
_ArrayLikeString_co,
_ArrayLikeTD64_co,
Expand Down Expand Up @@ -247,8 +260,18 @@ _DTypeT_co = TypeVar("_DTypeT_co", bound=dtype, default=dtype, covariant=True)
_ArrayT = TypeVar("_ArrayT", bound=ndarray[Any, Any])
_ScalarT = TypeVar("_ScalarT", bound=generic)
_ScalarT_co = TypeVar("_ScalarT_co", bound=generic, covariant=True)
_NumberT = TypeVar("_NumberT", bound=number)
# A subset of `MaskedArray` that can be parametrized w.r.t. `np.generic`
_MaskedArray: TypeAlias = MaskedArray[_AnyShape, dtype[_ScalarT]]

_MaskedArrayUInt_co: TypeAlias = _MaskedArray[unsignedinteger | np.bool]
_MaskedArrayInt_co: TypeAlias = _MaskedArray[integer | np.bool]
_MaskedArrayComplex_co: TypeAlias = _MaskedArray[inexact | integer | np.bool]
_MaskedArrayTD64_co: TypeAlias = _MaskedArray[timedelta64 | integer | np.bool]
_MaskedArrayFloat64_co: TypeAlias = _MaskedArray[floating[_64Bit] | float32 | float16 | integer | np.bool]
_MaskedArrayComplex128_co: TypeAlias = _MaskedArray[number[_64Bit] | number[_32Bit] | float16 | integer | np.bool]
_MaskedArrayFloat_co: TypeAlias = _MaskedArray[floating | integer | np.bool]

_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_ScalarT]]

MaskType = bool_
Expand Down Expand Up @@ -463,10 +486,179 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
def __gt__(self, other: ArrayLike, /) -> _MaskedArray[bool_]: ... # type: ignore[override]
def __le__(self, other: ArrayLike, /) -> _MaskedArray[bool_]: ... # type: ignore[override]
def __lt__(self, other: ArrayLike, /) -> _MaskedArray[bool_]: ... # type: ignore[override]
def __add__(self, other): ...
def __radd__(self, other): ...
def __sub__(self, other): ...
def __rsub__(self, other): ...

# Keep in sync with `ndarray.__add__`
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it might also be a good idea to mirror these comments on ndarray side in __init__.pyi then

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would you mind adding those there?

Copy link
Member Author

@MarcoGorelli MarcoGorelli May 27, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure what I'm missing sorry 😄

in ndarray.__add__, I don't see comments:

numpy/numpy/__init__.pyi

Lines 2811 to 2857 in 4125f0a

@overload
def __add__(self: NDArray[_NumberT], other: _ArrayLikeBool_co, /) -> NDArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: NDArray[np.bool], other: _ArrayLikeBool_co, /) -> NDArray[np.bool]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: NDArray[np.bool], other: _ArrayLike[_NumberT], /) -> NDArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: NDArray[float64], other: _ArrayLikeFloat64_co, /) -> NDArray[float64]: ...
@overload
def __add__(self: _ArrayFloat64_co, other: _ArrayLike[floating[_64Bit]], /) -> NDArray[float64]: ...
@overload
def __add__(self: NDArray[complex128], other: _ArrayLikeComplex128_co, /) -> NDArray[complex128]: ...
@overload
def __add__(self: _ArrayComplex128_co, other: _ArrayLike[complexfloating[_64Bit]], /) -> NDArray[complex128]: ...
@overload
def __add__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co, /) -> NDArray[unsignedinteger]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _ArrayInt_co, other: _ArrayLikeInt_co, /) -> NDArray[signedinteger]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co, /) -> NDArray[floating]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co, /) -> NDArray[complexfloating]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: NDArray[number], other: _ArrayLikeNumber_co, /) -> NDArray[number]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co, /) -> NDArray[timedelta64]: ...
@overload
def __add__(self: _ArrayTD64_co, other: _ArrayLikeDT64_co, /) -> NDArray[datetime64]: ...
@overload
def __add__(self: NDArray[datetime64], other: _ArrayLikeTD64_co, /) -> NDArray[datetime64]: ...
@overload
def __add__(self: NDArray[bytes_], other: _ArrayLikeBytes_co, /) -> NDArray[bytes_]: ...
@overload
def __add__(self: NDArray[str_], other: _ArrayLikeStr_co, /) -> NDArray[str_]: ...
@overload
def __add__(
self: ndarray[Any, dtypes.StringDType],
other: _ArrayLikeStr_co | _ArrayLikeString_co,
/,
) -> ndarray[tuple[Any, ...], dtypes.StringDType]: ...
@overload
def __add__(self: NDArray[object_], other: Any, /) -> Any: ...
@overload
def __add__(self: NDArray[Any], other: _ArrayLikeObject_co, /) -> Any: ...
@overload # signature equivalent to __add__
def __radd__(self: NDArray[_NumberT], other: int | np.bool, /) -> ndarray[_ShapeT_co, dtype[_NumberT]]: ...

in ndarray.__radd__, there's a single comment, which I've mirrored:

numpy/numpy/__init__.pyi

Line 2856 in 4125f0a

@overload # signature equivalent to __add__

and in ndarray.__sub__ and ndarray.__rsub__ i don't see comments

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh haha. I meant that it would help to add it to __init__.pyi as well, because it indeed isn't there yet :)

@overload
def __add__(self: _MaskedArray[_NumberT], other: int | np.bool, /) -> MaskedArray[_ShapeT_co, dtype[_NumberT]]: ...
@overload
def __add__(self: _MaskedArray[_NumberT], other: _ArrayLikeBool_co, /) -> _MaskedArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _MaskedArray[np.bool], other: _ArrayLikeBool_co, /) -> _MaskedArray[np.bool]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _MaskedArray[np.bool], other: _ArrayLike[_NumberT], /) -> _MaskedArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _MaskedArray[float64], other: _ArrayLikeFloat64_co, /) -> _MaskedArray[float64]: ...
@overload
def __add__(self: _MaskedArrayFloat64_co, other: _ArrayLike[floating[_64Bit]], /) -> _MaskedArray[float64]: ...
@overload
def __add__(self: _MaskedArray[complex128], other: _ArrayLikeComplex128_co, /) -> _MaskedArray[complex128]: ...
@overload
def __add__(self: _MaskedArrayComplex128_co, other: _ArrayLike[complexfloating[_64Bit]], /) -> _MaskedArray[complex128]: ...
@overload
def __add__(self: _MaskedArrayUInt_co, other: _ArrayLikeUInt_co, /) -> _MaskedArray[unsignedinteger]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _MaskedArrayInt_co, other: _ArrayLikeInt_co, /) -> _MaskedArray[signedinteger]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _MaskedArrayFloat_co, other: _ArrayLikeFloat_co, /) -> _MaskedArray[floating]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _MaskedArrayComplex_co, other: _ArrayLikeComplex_co, /) -> _MaskedArray[complexfloating]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _MaskedArray[number], other: _ArrayLikeNumber_co, /) -> _MaskedArray[number]: ... # type: ignore[overload-overlap]
@overload
def __add__(self: _MaskedArrayTD64_co, other: _ArrayLikeTD64_co, /) -> _MaskedArray[timedelta64]: ...
@overload
def __add__(self: _MaskedArrayTD64_co, other: _ArrayLikeDT64_co, /) -> _MaskedArray[datetime64]: ...
@overload
def __add__(self: _MaskedArray[datetime64], other: _ArrayLikeTD64_co, /) -> _MaskedArray[datetime64]: ...
@overload
def __add__(self: _MaskedArray[bytes_], other: _ArrayLikeBytes_co, /) -> _MaskedArray[bytes_]: ...
@overload
def __add__(self: _MaskedArray[str_], other: _ArrayLikeStr_co, /) -> _MaskedArray[str_]: ...
@overload
def __add__(
self: MaskedArray[Any, dtypes.StringDType],
other: _ArrayLikeStr_co | _ArrayLikeString_co,
/,
) -> MaskedArray[tuple[Any, ...], dtypes.StringDType]: ...
@overload
def __add__(self: _MaskedArray[object_], other: Any, /) -> Any: ...
@overload
def __add__(self: _MaskedArray[Any], other: _ArrayLikeObject_co, /) -> Any: ...

# Keep in sync with `ndarray.__radd__`
@overload
def __radd__(self: _MaskedArray[_NumberT], other: int | np.bool, /) -> MaskedArray[_ShapeT_co, dtype[_NumberT]]: ...
@overload
def __radd__(self: _MaskedArray[_NumberT], other: _ArrayLikeBool_co, /) -> _MaskedArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __radd__(self: _MaskedArray[np.bool], other: _ArrayLikeBool_co, /) -> _MaskedArray[np.bool]: ... # type: ignore[overload-overlap]
@overload
def __radd__(self: _MaskedArray[np.bool], other: _ArrayLike[_NumberT], /) -> _MaskedArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __radd__(self: _MaskedArray[float64], other: _ArrayLikeFloat64_co, /) -> _MaskedArray[float64]: ...
@overload
def __radd__(self: _MaskedArrayFloat64_co, other: _ArrayLike[floating[_64Bit]], /) -> _MaskedArray[float64]: ...
@overload
def __radd__(self: _MaskedArray[complex128], other: _ArrayLikeComplex128_co, /) -> _MaskedArray[complex128]: ...
@overload
def __radd__(self: _MaskedArrayComplex128_co, other: _ArrayLike[complexfloating[_64Bit]], /) -> _MaskedArray[complex128]: ...
@overload
def __radd__(self: _MaskedArrayUInt_co, other: _ArrayLikeUInt_co, /) -> _MaskedArray[unsignedinteger]: ... # type: ignore[overload-overlap]
@overload
def __radd__(self: _MaskedArrayInt_co, other: _ArrayLikeInt_co, /) -> _MaskedArray[signedinteger]: ... # type: ignore[overload-overlap]
@overload
def __radd__(self: _MaskedArrayFloat_co, other: _ArrayLikeFloat_co, /) -> _MaskedArray[floating]: ... # type: ignore[overload-overlap]
@overload
def __radd__(self: _MaskedArrayComplex_co, other: _ArrayLikeComplex_co, /) -> _MaskedArray[complexfloating]: ... # type: ignore[overload-overlap]
@overload
def __radd__(self: _MaskedArray[number], other: _ArrayLikeNumber_co, /) -> _MaskedArray[number]: ... # type: ignore[overload-overlap]
@overload
def __radd__(self: _MaskedArrayTD64_co, other: _ArrayLikeTD64_co, /) -> _MaskedArray[timedelta64]: ...
@overload
def __radd__(self: _MaskedArrayTD64_co, other: _ArrayLikeDT64_co, /) -> _MaskedArray[datetime64]: ...
@overload
def __radd__(self: _MaskedArray[datetime64], other: _ArrayLikeTD64_co, /) -> _MaskedArray[datetime64]: ...
@overload
def __radd__(self: _MaskedArray[bytes_], other: _ArrayLikeBytes_co, /) -> _MaskedArray[bytes_]: ...
@overload
def __radd__(self: _MaskedArray[str_], other: _ArrayLikeStr_co, /) -> _MaskedArray[str_]: ...
@overload
def __radd__(
self: MaskedArray[Any, dtypes.StringDType],
other: _ArrayLikeStr_co | _ArrayLikeString_co,
/,
) -> MaskedArray[tuple[Any, ...], dtypes.StringDType]: ...
@overload
def __radd__(self: _MaskedArray[object_], other: Any, /) -> Any: ...
@overload
def __radd__(self: _MaskedArray[Any], other: _ArrayLikeObject_co, /) -> Any: ...

# Keep in sync with `ndarray.__sub__`
@overload
def __sub__(self: _MaskedArray[_NumberT], other: int | np.bool, /) -> MaskedArray[_ShapeT_co, dtype[_NumberT]]: ...
@overload
def __sub__(self: _MaskedArray[_NumberT], other: _ArrayLikeBool_co, /) -> _MaskedArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __sub__(self: _MaskedArray[np.bool], other: _ArrayLikeBool_co, /) -> NoReturn: ...
@overload
def __sub__(self: _MaskedArray[np.bool], other: _ArrayLike[_NumberT], /) -> _MaskedArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __sub__(self: _MaskedArray[float64], other: _ArrayLikeFloat64_co, /) -> _MaskedArray[float64]: ...
@overload
def __sub__(self: _MaskedArrayFloat64_co, other: _ArrayLike[floating[_64Bit]], /) -> _MaskedArray[float64]: ...
@overload
def __sub__(self: _MaskedArray[complex128], other: _ArrayLikeComplex128_co, /) -> _MaskedArray[complex128]: ...
@overload
def __sub__(self: _MaskedArrayComplex128_co, other: _ArrayLike[complexfloating[_64Bit]], /) -> _MaskedArray[complex128]: ...
@overload
def __sub__(self: _MaskedArrayUInt_co, other: _ArrayLikeUInt_co, /) -> _MaskedArray[unsignedinteger]: ... # type: ignore[overload-overlap]
@overload
def __sub__(self: _MaskedArrayInt_co, other: _ArrayLikeInt_co, /) -> _MaskedArray[signedinteger]: ... # type: ignore[overload-overlap]
@overload
def __sub__(self: _MaskedArrayFloat_co, other: _ArrayLikeFloat_co, /) -> _MaskedArray[floating]: ... # type: ignore[overload-overlap]
@overload
def __sub__(self: _MaskedArrayComplex_co, other: _ArrayLikeComplex_co, /) -> _MaskedArray[complexfloating]: ... # type: ignore[overload-overlap]
@overload
def __sub__(self: _MaskedArray[number], other: _ArrayLikeNumber_co, /) -> _MaskedArray[number]: ... # type: ignore[overload-overlap]
@overload
def __sub__(self: _MaskedArrayTD64_co, other: _ArrayLikeTD64_co, /) -> _MaskedArray[timedelta64]: ...
@overload
def __sub__(self: _MaskedArray[datetime64], other: _ArrayLikeTD64_co, /) -> _MaskedArray[datetime64]: ...
@overload
def __sub__(self: _MaskedArray[datetime64], other: _ArrayLikeDT64_co, /) -> _MaskedArray[timedelta64]: ...
@overload
def __sub__(self: _MaskedArray[object_], other: Any, /) -> Any: ...
@overload
def __sub__(self: _MaskedArray[Any], other: _ArrayLikeObject_co, /) -> Any: ...

# Keep in sync with `ndarray.__rsub__`
@overload
def __rsub__(self: _MaskedArray[_NumberT], other: int | np.bool, /) -> MaskedArray[_ShapeT_co, dtype[_NumberT]]: ...
@overload
def __rsub__(self: _MaskedArray[_NumberT], other: _ArrayLikeBool_co, /) -> _MaskedArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __rsub__(self: _MaskedArray[np.bool], other: _ArrayLikeBool_co, /) -> NoReturn: ...
@overload
def __rsub__(self: _MaskedArray[np.bool], other: _ArrayLike[_NumberT], /) -> _MaskedArray[_NumberT]: ... # type: ignore[overload-overlap]
@overload
def __rsub__(self: _MaskedArray[float64], other: _ArrayLikeFloat64_co, /) -> _MaskedArray[float64]: ...
@overload
def __rsub__(self: _MaskedArrayFloat64_co, other: _ArrayLike[floating[_64Bit]], /) -> _MaskedArray[float64]: ...
@overload
def __rsub__(self: _MaskedArray[complex128], other: _ArrayLikeComplex128_co, /) -> _MaskedArray[complex128]: ...
@overload
def __rsub__(self: _MaskedArrayComplex128_co, other: _ArrayLike[complexfloating[_64Bit]], /) -> _MaskedArray[complex128]: ...
@overload
def __rsub__(self: _MaskedArrayUInt_co, other: _ArrayLikeUInt_co, /) -> _MaskedArray[unsignedinteger]: ... # type: ignore[overload-overlap]
@overload
def __rsub__(self: _MaskedArrayInt_co, other: _ArrayLikeInt_co, /) -> _MaskedArray[signedinteger]: ... # type: ignore[overload-overlap]
@overload
def __rsub__(self: _MaskedArrayFloat_co, other: _ArrayLikeFloat_co, /) -> _MaskedArray[floating]: ... # type: ignore[overload-overlap]
@overload
def __rsub__(self: _MaskedArrayComplex_co, other: _ArrayLikeComplex_co, /) -> _MaskedArray[complexfloating]: ... # type: ignore[overload-overlap]
@overload
def __rsub__(self: _MaskedArray[number], other: _ArrayLikeNumber_co, /) -> _MaskedArray[number]: ... # type: ignore[overload-overlap]
@overload
def __rsub__(self: _MaskedArrayTD64_co, other: _ArrayLikeTD64_co, /) -> _MaskedArray[timedelta64]: ...
@overload
def __rsub__(self: _MaskedArrayTD64_co, other: _ArrayLikeDT64_co, /) -> _MaskedArray[datetime64]: ...
@overload
def __rsub__(self: _MaskedArray[datetime64], other: _ArrayLikeDT64_co, /) -> _MaskedArray[timedelta64]: ...
@overload
def __rsub__(self: _MaskedArray[object_], other: Any, /) -> Any: ...
@overload
def __rsub__(self: _MaskedArray[Any], other: _ArrayLikeObject_co, /) -> Any: ...

def __mul__(self, other): ...
def __rmul__(self, other): ...
def __truediv__(self, other): ...
Expand Down
Loading
Loading
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy