Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit d946482

Browse filesBrowse files
MarcoGorelliMarcoGorelli
committed
TYP: Type MaskedArray.__ge__ 
1 parent 8f11c7b commit d946482
Copy full SHA for d946482

File tree

Expand file treeCollapse file tree

3 files changed

+46
-4
lines changed
Filter options
Expand file treeCollapse file tree

3 files changed

+46
-4
lines changed

‎numpy/ma/core.pyi

Copy file name to clipboardExpand all lines: numpy/ma/core.pyi
+27-4Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,22 +15,31 @@ from numpy import (
1515
amax,
1616
amin,
1717
bool_,
18+
datetime64,
1819
dtype,
1920
expand_dims,
2021
float64,
2122
generic,
23+
object_,
24+
integer,
2225
intp,
2326
ndarray,
27+
number,
28+
timedelta64,
2429
)
2530
from numpy._globals import _NoValueType
2631
from numpy._typing import (
2732
ArrayLike,
28-
_IntLike_co,
2933
NDArray,
3034
_ArrayLike,
35+
_ArrayLikeInt,
3136
_ArrayLikeInt_co,
37+
_ArrayLikeNumber_co,
38+
_ArrayLikeDT64_co,
39+
_ArrayLikeObject_co,
40+
_ArrayLikeTD64_co,
3241
_DTypeLikeBool,
33-
_ArrayLikeInt,
42+
_IntLike_co,
3443
_ScalarLike_co,
3544
_Shape,
3645
_ShapeLike,
@@ -223,8 +232,11 @@ _DType = TypeVar("_DType", bound=dtype[Any])
223232
_DType_co = TypeVar("_DType_co", bound=dtype[Any], covariant=True)
224233
_ArrayT = TypeVar("_ArrayT", bound=ndarray[Any, Any])
225234
_SCT = TypeVar("_SCT", bound=generic)
235+
_SCT_co = TypeVar("_SCT_co", bound=generic, covariant=True)
226236
# A subset of `MaskedArray` that can be parametrized w.r.t. `np.generic`
227-
_MaskedArray: TypeAlias = MaskedArray[_Shape, dtype[_SCT]]
237+
_MaskedArray: TypeAlias = MaskedArray[_Shape, dtype[_SCT_co]]
238+
_MaskedArrayNumber_co: TypeAlias = _MaskedArray[number | bool_]
239+
_MaskedArrayTD64_co: TypeAlias = _MaskedArray[timedelta64 | integer | bool_]
228240

229241
MaskType = bool
230242
nomask: bool
@@ -422,7 +434,18 @@ class MaskedArray(ndarray[_ShapeT_co, _DType_co]):
422434
def compress(self, condition, axis=..., out=...): ...
423435
def __eq__(self, other): ...
424436
def __ne__(self, other): ...
425-
def __ge__(self, other): ...
437+
438+
@overload
439+
def __ge__(self: _MaskedArrayNumber_co, other: _ArrayLikeNumber_co, /) -> _MaskedArray[bool_]: ...
440+
@overload
441+
def __ge__(self: _MaskedArrayTD64_co, other: _ArrayLikeTD64_co, /) -> _MaskedArray[bool_]: ...
442+
@overload
443+
def __ge__(self: _MaskedArray[datetime64], other: _ArrayLikeDT64_co, /) -> _MaskedArray[bool_]: ...
444+
@overload
445+
def __ge__(self: _MaskedArray[object_], other: Any, /) -> _MaskedArray[bool_]: ...
446+
@overload
447+
def __ge__(self: _MaskedArray[Any], other: _ArrayLikeObject_co, /) -> _MaskedArray[bool_]: ...
448+
426449
def __gt__(self, other): ...
427450
def __le__(self, other): ...
428451
def __lt__(self, other): ...

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

Copy file name to clipboardExpand all lines: numpy/typing/tests/data/fail/ma.pyi
+2Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,3 +92,5 @@ m.argpartition(axis=(0,1)) # E: No overload variant
9292
m.argpartition(kind='cabbage') # E: No overload variant
9393
m.argpartition(order=lambda: 'cabbage') # E: No overload variant
9494
m.argpartition(AR_b) # E: No overload variant
95+
96+
m >= 'mango' # E: No overload variant

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

Copy file name to clipboardExpand all lines: numpy/typing/tests/data/reveal/ma.pyi
+17Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ from typing_extensions import assert_type
33
from typing import Any, TypeAlias, TypeVar
44
from numpy._typing import _Shape, NDArray
55
from numpy import dtype, generic
6+
from datetime import datetime, timedelta
67

78

89
_SCT_co = TypeVar("_SCT_co", bound=generic, covariant=True)
@@ -11,12 +12,20 @@ MaskedNDArray: TypeAlias = np.ma.MaskedArray[_Shape, dtype[_SCT_co]]
1112
class MaskedNDArraySubclass(MaskedNDArray[np.complex128]): ...
1213

1314
AR_f4: NDArray[np.float32]
15+
AR_dt64: NDArray[np.datetime64]
16+
AR_td64: NDArray[np.timedelta64]
17+
AR_o: NDArray[np.timedelta64]
1418

1519
MAR_b: MaskedNDArray[np.bool]
1620
MAR_f4: MaskedNDArray[np.float32]
1721
MAR_f8: MaskedNDArray[np.float64]
1822
MAR_i8: MaskedNDArray[np.int64]
23+
MAR_o: MaskedNDArray[np.object_]
24+
MAR_dt64: MaskedNDArray[np.datetime64]
25+
MAR_td64: MaskedNDArray[np.timedelta64]
26+
1927
MAR_subclass: MaskedNDArraySubclass
28+
2029
MAR_1d: np.ma.MaskedArray[tuple[int], np.dtype[Any]]
2130

2231
f4: np.float32
@@ -156,3 +165,11 @@ assert_type(MAR_f4.partition(1, axis=0, kind='introselect', order='K'), None)
156165

157166
assert_type(MAR_f4.argpartition(1), MaskedNDArray[np.intp])
158167
assert_type(MAR_1d.argpartition(1, axis=0, kind='introselect', order='K'), MaskedNDArray[np.intp])
168+
169+
assert_type(MAR_f4 >= 3, MaskedNDArray[np.bool])
170+
assert_type(MAR_i8 >= AR_td64, MaskedNDArray[np.bool])
171+
assert_type(MAR_b >= AR_td64, MaskedNDArray[np.bool])
172+
assert_type(MAR_td64 >= AR_td64, MaskedNDArray[np.bool])
173+
assert_type(MAR_dt64 >= AR_dt64, MaskedNDArray[np.bool])
174+
assert_type(MAR_o >= AR_o, MaskedNDArray[np.bool])
175+
assert_type(MAR_1d >= 0, MaskedNDArray[np.bool])

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.