8000 TYP: Type ``MaskedArray.{imag, real, baseclass, mT}`` (#28868) · numpy/numpy@42205e0 · GitHub
[go: up one dir, main page]

Skip to content

Commit 42205e0

Browse files
authored
TYP: Type MaskedArray.{imag, real, baseclass, mT} (#28868)
1 parent 1752876 commit 42205e0

File tree

2 files changed

+19
-11
lines changed

2 files changed

+19
-11
lines changed

numpy/ma/core.pyi

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,14 @@
22
# ruff: noqa: ANN001, ANN002, ANN003, ANN201, ANN202 ANN204, ANN401
33

44
from collections.abc import Sequence
5-
from typing import Any, Literal, SupportsIndex, TypeAlias, TypeVar, overload
5+
from typing import Any, Literal, Self, SupportsIndex, TypeAlias, TypeVar, overload
66

77
from _typeshed import Incomplete
88
from typing_extensions import TypeIs, deprecated
99

1010
import numpy as np
1111
from numpy import (
12+
_HasDTypeWithRealAndImag,
1213
_ModeKind,
1314
_OrderKACF,
1415
_PartitionKind,
@@ -401,10 +402,6 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
401402
def __getitem__(self, indx): ...
402403
def __setitem__(self, indx, value): ...
403404
@property
404-
def dtype(self) -> _DTypeT_co: ...
405-
@dtype.setter
406-
def dtype(self: MaskedArray[Any, _DTypeT], dtype: _DTypeT, /) -> None: ...
407-
@property
408405
def shape(self) -> _ShapeT_co: ...
409406
@shape.setter
410407
def shape(self: MaskedArray[_ShapeT, Any], shape: _ShapeT, /) -> None: ...
@@ -426,7 +423,7 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
426423
def sharedmask(self): ...
427424
def shrink_mask(self): ...
428425
@property
429-
def baseclass(self): ...
426+
def baseclass(self) -> type[NDArray[Any]]: ...
430427
data: Any
431428
@property
432429
def flat(self): ...
@@ -468,10 +465,10 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
468465
def __itruediv__(self, other): ...
469466
def __ipow__(self, other): ...
470467
@property # type: ignore[misc]
471-
def imag(self): ...
468+
def imag(self: _HasDTypeWithRealAndImag[object, _ScalarT], /) -> MaskedArray[_ShapeT_co, dtype[_ScalarT]]: ...
472469
get_imag: Any
473470
@property # type: ignore[misc]
474-
def real(self): ...
471+
def real(self: _HasDTypeWithRealAndImag[_ScalarT, object], /) -> MaskedArray[_ShapeT_co, dtype[_ScalarT]]: ...
475472
get_real: Any
476473

477474
# keep in sync with `np.ma.count`
@@ -804,9 +801,6 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
804801
T: Any
805802
transpose: Any
806803

807-
@property # type: ignore[misc]
808-
def mT(self): ...
809-
810804
#
811805
def toflex(self) -> Incomplete: ...
812806
def torecords(self) -> Incomplete: ...
@@ -820,6 +814,12 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
820814
def __reduce__(self): ...
821815
def __deepcopy__(self, memo=...): ...
822816

817+
# Keep `dtype` at the bottom to avoid name conflicts with `np.dtype`
818+
@property
819+
def dtype(self) -> _DTypeT_co: ...
820< 8000 code class="diff-text syntax-highlighted-line addition">+
@dtype.setter
821+
def dtype(self: MaskedArray[Any, _DTypeT], dtype: _DTypeT, /) -> None: ...
822+
823823
class mvoid(MaskedArray[_ShapeT_co, _DTypeT_co]):
824824
def __new__(
825825
self, # pyright: ignore[reportSelfClsParameterName]

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ AR_dt64: NDArray[np.datetime64]
1717
AR_td64: NDArray[np.timedelta64]
1818
AR_o: NDArray[np.timedelta64]
1919

20+
MAR_c16: MaskedArray[np.complex128]
2021
MAR_b: MaskedArray[np.bool]
2122
MAR_f4: MaskedArray[np.float32]
2223
MAR_f8: MaskedArray[np.float64]
@@ -309,6 +310,13 @@ def func(x: object) -> None:
309310
else:
310311
assert_type(x, object)
311312

313+
assert_type(MAR_2d_f4.mT, np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
314+
315+
assert_type(MAR_c16.real, MaskedArray[np.float64])
316+
assert_type(MAR_c16.imag, MaskedArray[np.float64])
317+
318+
assert_type(MAR_2d_f4.baseclass, type[NDArray[Any]])
319+
312320
assert_type(MAR_b.swapaxes(0, 1), MaskedArray[np.bool])
313321
assert_type(MAR_2d_f4.swapaxes(1, 0), MaskedArray[np.float32])
314322

0 commit comments

Comments
 (0)
0