🌐 AI搜索 & 代理 主页
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

Closed
MarcoGorelli wants to merge 5 commits into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c23500e
TYP: Type ``np.ma.squeeze`` and ``np.ma.{mask_rows, mask_cols, mask_r…
MarcoGorelli Apr 21, 2025
656ac92
reorder imports
MarcoGorelli Apr 22, 2025
40fdc0a
specialise return shape to tuple[int, int]
MarcoGorelli Apr 22, 2025
5ab323d
imports...
MarcoGorelli Apr 22, 2025
3fe03e5
remove unused imports
MarcoGorelli Apr 22, 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
19 changes: 17 additions & 2 deletions numpy/ma/core.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1205,10 +1205,25 @@ identity: _convert2ma
indices: _convert2ma
ones: _convert2ma
ones_like: _convert2ma
squeeze: _convert2ma

@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]: ...
Comment on lines +1209 to +1223
Copy link
Member

@jorenham jorenham Apr 23, 2025

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?


zeros: _convert2ma
zeros_like: _convert2ma

def append(a, b, axis=...): ...
def dot(a, b, strict=..., out=...): ...
def mask_rowcols(a, axis=...): ...
29 changes: 24 additions & 5 deletions numpy/ma/extras.pyi
View file
Open in desktop
Copy link
Member

@jorenham jorenham Apr 23, 2025

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.

Original file line number Diff line number Diff line change
@@ -1,10 +1,18 @@
from typing import Any, SupportsIndex, overload

from _typeshed import Incomplete

import numpy as np
from numpy.lib._function_base_impl import average
from numpy.lib._index_tricks_impl import AxisConcatenator
from numpy.typing import ArrayLike
from numpy._typing import _ArrayLike
Comment on lines +8 to +9
Copy link
Member

@jorenham jorenham Apr 23, 2025

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


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

@jorenham jorenham Apr 23, 2025

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


__all__ = [
"apply_along_axis",
Expand Down Expand Up @@ -96,8 +104,17 @@ def compress_nd(x, axis=...): ...
def compress_rowcols(x, axis=...): ...
def compress_rows(a): ...
def compress_cols(a): ...
def mask_rows(a, axis = ...): ...
def mask_cols(a, axis = ...): ...

@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

@jorenham jorenham Apr 23, 2025

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]: ...


@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]]: ...
Comment on lines +108 to +116
Copy link
Member Author

@MarcoGorelli MarcoGorelli Apr 22, 2025

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

@MarcoGorelli MarcoGorelli Apr 22, 2025

Choose a reason for hiding this comment

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

*stubs

Copy link
Member

@jorenham jorenham Apr 23, 2025

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

@MarcoGorelli MarcoGorelli Apr 26, 2025

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.


def ediff1d(arr, to_end=..., to_begin=...): ...
def unique(ar1, return_index=..., return_inverse=...): ...
def intersect1d(ar1, ar2, assume_unique=...): ...
Expand Down Expand Up @@ -130,5 +147,7 @@ def clump_masked(a): ...
def vander(x, n=...): ...
def polyfit(x, y, deg, rcond=..., full=..., w=..., cov=...): ...

#
def mask_rowcols(a: Incomplete, axis: Incomplete | None = None) -> MaskedArray[Incomplete, np.dtype[Incomplete]]: ...
@overload
def mask_rowcols(a: _ArrayLike[_ScalarT_co], axis: SupportsIndex | None = None) -> MaskedArray[tuple[int, int], np.dtype[_ScalarT_co]]: ...
@overload
def mask_rowcols(a: ArrayLike, axis: SupportsIndex | None = None) -> MaskedArray[tuple[int, int], np.dtype]: ...
10 changes: 9 additions & 1 deletion numpy/typing/tests/data/fail/ma.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ import numpy.ma
import numpy.typing as npt

m: np.ma.MaskedArray[tuple[int], np.dtype[np.float64]]

AR_b: npt.NDArray[np.bool]
MAR_2d_f4: np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]]

m.shape = (3, 1) # E: Incompatible types in assignment
m.dtype = np.bool # E: Incompatible types in assignment
Expand Down Expand Up @@ -115,4 +115,12 @@ np.ma.put(m, 4, 999, mode='flip') # E: No overload variant

np.ma.put([1,1,3], 0, 999) # E: No overload variant

np.ma.squeeze(m, 1.0) # E: No overload variant

np.ma.mask_rows(MAR_2d_f4, axis=0) # E: No overload variant

np.ma.mask_cols(MAR_2d_f4, axis=1) # E: No overload variant

np.ma.mask_rowcols(MAR_2d_f4, axis='broccoli') # E: No overload variant

np.ma.compressed(lambda: 'compress me') # E: No overload variant
28 changes: 28 additions & 0 deletions numpy/typing/tests/data/reveal/ma.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@jorenham jorenham Apr 23, 2025

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]

MAR_2d_f4: np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]]

b: np.bool
f4: np.float32
Expand Down Expand Up @@ -264,12 +266,38 @@ assert_type(np.ma.put(MAR_f4, 4, 999, mode='clip'), None)

assert_type(np.ma.putmask(MAR_f4, [True, False], [0, 1]), None)

assert_type(np.ma.squeeze(b), np.bool)
assert_type(np.ma.squeeze(f4), np.float32)
assert_type(np.ma.squeeze(f), MaskedNDArray[Any])
assert_type(np.ma.squeeze(MAR_b), MaskedNDArray[np.bool])
assert_type(np.ma.squeeze(AR_f4), MaskedNDArray[np.float32])

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

@jorenham jorenham Apr 23, 2025

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.

assert_type(np.ma.mask_rows(MAR_2d), np.ma.MaskedArray[tuple[int, int], np.dtype]) # type: ignore[assert-type]
Copy link
Member

@jorenham jorenham Apr 23, 2025

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?


assert_type(np.ma.mask_cols(MAR_2d_f4), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.ma.mask_cols(MAR_f4), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.ma.mask_cols([[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_cols(MAR_2d), np.ma.MaskedArray[tuple[int, int], np.dtype]) # type: ignore[assert-type]

assert_type(np.ma.mask_rowcols(MAR_2d_f4), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.ma.mask_rowcols(MAR_f4), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.ma.mask_rowcols(MAR_2d_f4, axis=0), np.ma.MaskedArray[tuple[int, int], np.dtype[np.float32]])
assert_type(np.ma.mask_rowcols([[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_rowcols(MAR_2d), np.ma.MaskedArray[tuple[int, int], np.dtype]) # type: ignore[assert-type]

assert_type(MAR_f4.filled(float('nan')), NDArray[np.float32])
assert_type(MAR_i8.filled(), NDArray[np.int64])
assert_type(MAR_1d.filled(), np.ndarray[tuple[int], np.dtype])

assert_type(np.ma.filled(MAR_f4, float('nan')), NDArray[np.float32])
assert_type(np.ma.filled([[1,2,3]]), NDArray[Any])
assert_type(np.ma.filled(MAR_2d_f4), np.ndarray[tuple[int, int], np.dtype[np.float32]])
# PyRight detects this one correctly, but mypy doesn't.
# https://github.com/numpy/numpy/pull/28742#discussion_r2048968375
assert_type(np.ma.filled(MAR_1d), np.ndarray[tuple[int], np.dtype]) # type: ignore[assert-type]
Loading