Thanks to visit codestin.com
Credit goes to github.com

Skip to content

TYP: Type np.ma.squeeze and np.ma.{mask_rows, mask_cols, mask_rowcols} #28794

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

Jump to bottom
Draft
wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

TYP: Type np.ma.squeeze and np.ma.{mask_rows, mask_cols, mask_rowcols} #28794

wants to merge 5 commits into from
+78 −8

Conversation

MarcoGorelli
Copy link
Member

No description provided.

MarcoGorelli
MarcoGorelli commented Apr 22, 2025
View reviewed changes
numpy/ma/extras.pyi
Comment on lines +111 to +119
@overload
def mask_rows(a: _ArrayLike[_ScalarT_co]) -> MaskedArray[tuple[int, int], np.dtype[_ScalarT_co]]: ...
@overload
def mask_rows(a: ArrayLike) -> MaskedArray[tuple[int, int], np.dtype[Any]]: ...

@overload
def mask_cols(a: _ArrayLike[_ScalarT_co]) -> MaskedArray[tuple[int, int], np.dtype[_ScalarT_co]]: ...
@overload
def mask_cols(a: ArrayLike) -> MaskedArray[tuple[int, int], np.dtype[Any]]: ...
Copy link
Member Author

Choose a reason for hiding this comment

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

technically these both accept a axis argument, but it emits a deprecationwarning, so i've left it out of the stops so as to discourage its usage

Copy link
Member Author

Choose a reason for hiding this comment

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

*stubs

Copy link
Member

Choose a reason for hiding this comment

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

Hmm stubtest won't like that. What I usually do in these situations, is add another @overload with the axis kwarg that has no default, and mark it as @deprecated. See e.g.

numpy/numpy/_core/arrayprint.pyi

Lines 88 to 180 in 1bc8764

# public numpy export
@overload # no style
def array2string(
a: NDArray[Any],
max_line_width: int | None = None,
precision: SupportsIndex | None = None,
suppress_small: bool | None = None,
separator: str = " ",
prefix: str = "",
style: _NoValueType = ...,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
*,
legacy: _Legacy | None = None,
) -> str: ...
@overload # style=<given> (positional), legacy="1.13"
def array2string(
a: NDArray[Any],
max_line_width: int | None,
precision: SupportsIndex | None,
suppress_small: bool | None,
separator: str,
prefix: str,
style: _ReprFunc,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
*,
legacy: Literal["1.13"],
) -> str: ...
@overload # style=<given> (keyword), legacy="1.13"
def array2string(
a: NDArray[Any],
max_line_width: int | None = None,
precision: SupportsIndex | None = None,
suppress_small: bool | None = None,
separator: str = " ",
prefix: str = "",
*,
style: _ReprFunc,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
legacy: Literal["1.13"],
) -> str: ...
@overload # style=<given> (positional), legacy!="1.13"
@deprecated("'style' argument is deprecated and no longer functional except in 1.13 'legacy' mode")
def array2string(
a: NDArray[Any],
max_line_width: int | None,
precision: SupportsIndex | None,
suppress_small: bool | None,
separator: str,
prefix: str,
style: _ReprFunc,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
*,
legacy: _LegacyNoStyle | None = None,
) -> str: ...
@overload # style=<given> (keyword), legacy="1.13"
@deprecated("'style' argument is deprecated and no longer functional except in 1.13 'legacy' mode")
def array2string(
a: NDArray[Any],
max_line_width: int | None = None,
precision: SupportsIndex | None = None,
suppress_small: bool | None = None,
separator: str = " ",
prefix: str = "",
*,
style: _ReprFunc,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
legacy: _LegacyNoStyle | None = None,
) -> str: ...

Copy link
Member Author

Choose a reason for hiding this comment

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

thanks - how do you run stubtest?

I've tried

 spin run python -m mypy.stubtest numpy

but it seems to hang indefinitely. couldn't find anything else about it in the numpy repo

Copy link
Member

@jorenham jorenham Apr 26, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

At the moment stubtest is only used by numtype. If you'd try to run it on numpy's bundled stubs, you'd get thousands of (accurate) errors. But properly fixing those would result in many breaking changes. And that's one of the reasons for building numtype, actually. The plan is to eventually also be able to use stubtest in numpy, once the other features that I'm developing there (like shape-typing) are functional and sufficiently stable.

And as you've probably noted by now, I've been tracking your np.ma stubs changes, (numpy/numtype#456) with the intention of porting them over. So with that in mind, I'm trying to anticipate any stubtest errors. Because if those turn up when porting a numpy PR to numtype, then I'll eventually have to port those back from numtype to numpy again.

So to get the stubtest output, you could copy these changes to numtype (https://github.com/numpy/numtype/tree/main/src/numpy-stubs/ma), and then uv run tool/stubtest.py (see https://numpy.org/numtype/dev/ for details). It's also an option to first submit your ma PR's there, and then port them over to numpy. It sounds pretty inefficient, but it ensures that that stubtest, mypy, and pyright (configured in the strictest of modes), are also happy.
And just to be clear; submitting your pr's to numpy is no problem as far as I'm concerned. It's just that we're a bit more in the dark here, CI-wise.

@jorenham jorenham self-requested a review April 22, 2025 13:31
@MarcoGorelli MarcoGorelli marked this pull request as ready for review April 22, 2025 14:11
@jorenham jorenham added the component: numpy.ma masked arrays label Apr 23, 2025
@jorenham jorenham mentioned this pull request Apr 23, 2025
Open
30 tasks
jorenham
jorenham reviewed Apr 23, 2025
View reviewed changes
numpy/ma/extras.pyi
Comment on lines +111 to +119
@overload
def mask_rows(a: _ArrayLike[_ScalarT_co]) -> MaskedArray[tuple[int, int], np.dtype[_ScalarT_co]]: ...
@overload
def mask_rows(a: ArrayLike) -> MaskedArray[tuple[int, int], np.dtype[Any]]: ...

@overload
def mask_cols(a: _ArrayLike[_ScalarT_co]) -> MaskedArray[tuple[int, int], np.dtype[_ScalarT_co]]: ...
@overload
def mask_cols(a: ArrayLike) -> MaskedArray[tuple[int, int], np.dtype[Any]]: ...
Copy link
Member

Choose a reason for hiding this comment

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

Hmm stubtest won't like that. What I usually do in these situations, is add another @overload with the axis kwarg that has no default, and mark it as @deprecated. See e.g.

numpy/numpy/_core/arrayprint.pyi

Lines 88 to 180 in 1bc8764

# public numpy export
@overload # no style
def array2string(
a: NDArray[Any],
max_line_width: int | None = None,
precision: SupportsIndex | None = None,
suppress_small: bool | None = None,
separator: str = " ",
prefix: str = "",
style: _NoValueType = ...,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
*,
legacy: _Legacy | None = None,
) -> str: ...
@overload # style=<given> (positional), legacy="1.13"
def array2string(
a: NDArray[Any],
max_line_width: int | None,
precision: SupportsIndex | None,
suppress_small: bool | None,
separator: str,
prefix: str,
style: _ReprFunc,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
*,
legacy: Literal["1.13"],
) -> str: ...
@overload # style=<given> (keyword), legacy="1.13"
def array2string(
a: NDArray[Any],
max_line_width: int | None = None,
precision: SupportsIndex | None = None,
suppress_small: bool | None = None,
separator: str = " ",
prefix: str = "",
*,
style: _ReprFunc,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
legacy: Literal["1.13"],
) -> str: ...
@overload # style=<given> (positional), legacy!="1.13"
@deprecated("'style' argument is deprecated and no longer functional except in 1.13 'legacy' mode")
def array2string(
a: NDArray[Any],
max_line_width: int | None,
precision: SupportsIndex | None,
suppress_small: bool | None,
separator: str,
prefix: str,
style: _ReprFunc,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
*,
legacy: _LegacyNoStyle | None = None,
) -> str: ...
@overload # style=<given> (keyword), legacy="1.13"
@deprecated("'style' argument is deprecated and no longer functional except in 1.13 'legacy' mode")
def array2string(
a: NDArray[Any],
max_line_width: int | None = None,
precision: SupportsIndex | None = None,
suppress_small: bool | None = None,
separator: str = " ",
prefix: str = "",
*,
style: _ReprFunc,
formatter: _FormatDict | None = None,
threshold: int | None = None,
edgeitems: int | None = None,
sign: _Sign | None = None,
floatmode: _FloatMode | None = None,
suffix: str = "",
legacy: _LegacyNoStyle | None = None,
) -> str: ...

numpy/ma/core.pyi
Comment on lines +1209 to +1223
@overload
def squeeze(
a: _ScalarT,
axis: _ShapeLike | None = ...,
) -> _ScalarT: ...
@overload
def squeeze(
a: _ArrayLike[_ScalarT],
axis: _ShapeLike | None = ...,
) -> _MaskedArray[_ScalarT]: ...
@overload
def squeeze(
a: ArrayLike,
axis: _ShapeLike | None = ...,
) -> _MaskedArray[Any]: ...
Copy link
Member

Choose a reason for hiding this comment

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

So it turns out that stubtest doesn't accept these being actual functions, and instead requires them to be instances of numpy.ma.core._convert2ma. You can see this in numpy/numtype#468. It's currently a WIP, but there I (attempted to) explain how the _convert2ma can be made generic on it's __call__ signature, so that a callable Protocol can be used to specify the signature. It's a bit messy, but I'm afraid that it's the only way that stubtest well be able to swallow it.

Are you willing to attempt doing this here as well, or would you prefer that we address this at a later point?

numpy/ma/extras.pyi
@overload
def mask_rows(a: _ArrayLike[_ScalarT_co]) -> MaskedArray[tuple[int, int], np.dtype[_ScalarT_co]]: ...
@overload
def mask_rows(a: ArrayLike) -> MaskedArray[tuple[int, int], np.dtype[Any]]: ...
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
def mask_rows(a: ArrayLike) -> MaskedArray[tuple[int, int], np.dtype[Any]]: ...
def mask_rows(a: ArrayLike) -> MaskedArray[tuple[int, int], np.dtype]: ...

numpy/ma/extras.pyi
Comment on lines +11 to +15
from .core import (
MaskedArray,
dot,
_ScalarT_co,
)
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
from .core import (
MaskedArray,
dot,
_ScalarT_co,
)
from .core import MaskedArray, dot, _ScalarT_co

numpy/ma/extras.pyi
Comment on lines +8 to +9
from numpy.typing import ArrayLike
from numpy._typing import _ArrayLike
Copy link
Member

Choose a reason for hiding this comment

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

nitpick: usually we also import ArrayLike from _typing internally

Suggested change
from numpy.typing import ArrayLike
from numpy._typing import _ArrayLike
from numpy._typing import ArrayLike, _ArrayLike

numpy/typing/tests/data/reveal/ma.pyi
@@ -29,6 +29,8 @@ MAR_V: MaskedNDArray[np.void]
MAR_subclass: MaskedNDArraySubclass

MAR_1d: np.ma.MaskedArray[tuple[int], np.dtype]
MAR_2d: np.ma.MaskedArray[tuple[int, int], np.dtype[Any]]
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
MAR_2d: np.ma.MaskedArray[tuple[int, int], np.dtype[Any]]
MAR_2d: np.ma.MaskedArray[tuple[int, int], np.dtype]

numpy/typing/tests/data/reveal/ma.pyi
assert_type(np.ma.mask_rows(MAR_f4), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.ma.mask_rows([[1,2,3]]), np.ma.MaskedArray[tuple[int, int], np.dtype])
# PyRight detects this one correctly, but mypy doesn't.
assert_type(np.ma.mask_rows(MAR_2d), np.ma.MaskedArray[tuple[int, int], np.dtype]) # type: ignore[assert-type]
Copy link
Member

Choose a reason for hiding this comment

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

what does mypy infer it as?

numpy/typing/tests/data/reveal/ma.pyi
assert_type(np.ma.mask_rows(MAR_2d_f4), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.ma.mask_rows(MAR_f4), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.ma.mask_rows([[1,2,3]]), np.ma.MaskedArray[tuple[int, int], np.dtype])
# PyRight detects this one correctly, but mypy doesn't.
Copy link
Member

Choose a reason for hiding this comment

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

At https://github.com/microsoft/pyright it's written as "Pyright" 🤷🏻

Suggested change
# PyRight detects this one correctly, but mypy doesn't.
# Pyright detects this one correctly, but mypy doesn't.

numpy/ma/extras.pyi
Copy link
Member

Choose a reason for hiding this comment

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

There are several -> MaskedArray[tuple[int, int], np.dtype[...]] annotations in this module now. A private type-alias could help make it more dry.

@jorenham
Copy link
Member

needs a rebase

This was referenced Apr 28, 2025
Merged
Merged
Merged
@MarcoGorelli MarcoGorelli marked this pull request as draft April 28, 2025 20:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@jorenham jorenham jorenham left review comments

Assignees
No one assigned
Labels
41 - Static typing component: numpy.ma masked arrays
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@MarcoGorelli @jorenham