8000 Merge pull request #28747 from jorenham/typing/partition-fixes · numpy/numpy@101621c · GitHub
[go: up one dir, main page]
Skip to content
8000

Commit 101621c

Browse files
charrischarris
authored
Merge pull request #28747 from jorenham/typing/partition-fixes
TYP: some ``[arg]partition`` fixes
2 parents c244462 + 4a470ec commit 101621c

File tree

4 files changed

+90
-29
lines changed

4 files changed

+90
-29
lines changed

‎numpy/__init__.pyi

Lines changed: 37 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2290,14 +2290,47 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
22902290
where: _ArrayLikeBool_co = True,
22912291
) -> _ArrayT: ...
22922292

2293+
#
2294+
@overload
2295+
def partition(
2296+
self,
2297+
/,
2298+
kth: _ArrayLikeInt,
2299+
axis: SupportsIndex = -1,
2300+
kind: _PartitionKind = "introselect"< 10000 /span>,
2301+
order: None = None,
2302+
) -> None: ...
2303+
@overload
2304+
def partition(
2305+
self: NDArray[void],
2306+
/,
2307+
kth: _ArrayLikeInt,
2308+
axis: SupportsIndex = -1,
2309+
kind: _PartitionKind = "introselect",
2310+
order: str | Sequence[str] | None = None,
2311+
) -> None: ...
2312+
2313+
#
2314+
@overload
22932315
def argpartition(
22942316
self,
2295-
kth: _ArrayLikeInt_co,
2296-
axis: None | SupportsIndex = ...,
2297-
kind: _PartitionKind = ...,
2298-
order: None | str | Sequence[str] = ...,
2317+
/,
2318+
kth: _ArrayLikeInt,
2319+
axis: SupportsIndex | None = -1,
2320+
kind: _PartitionKind = "introselect",
2321+
order: None = None,
2322+
) -> NDArray[intp]: ...
2323+
@overload
2324+
def argpartition(
2325+
self: NDArray[void],
2326+
/,
2327+
kth: _ArrayLikeInt,
2328+
axis: SupportsIndex | None = -1,
2329+
kind: _PartitionKind = "introselect",
2330+
order: str | Sequence[str] | None = None,
22992331
) -> NDArray[intp]: ...
23002332

2333+
#
23012334
def diagonal(
23022335
self,
23032336
offset: SupportsIndex = ...,
@@ -2317,14 +2350,6 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
23172350
# `nonzero()` is deprecated for 0d arrays/generics
23182351
def nonzero(self) -> tuple[NDArray[intp], ...]: ...
23192352

2320-
def partition(
2321-
self,
2322-
kth: _ArrayLikeInt_co,
2323-
axis: SupportsIndex = ...,
2324-
kind: _PartitionKind = ...,
2325-
order: None | str | Sequence[str] = ...,
2326-
) -> None: ...
2327-
23282353
# `put` is technically available to `generic`,
23292354
# but is pointless as `generic`s are immutable
23302355
def put(self, /, indices: _ArrayLikeInt_co, values: ArrayLike, mode: _ModeKind = "raise") -> None: ...

‎numpy/_core/fromnumeric.pyi

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ from typing_extensions import deprecated
1717

1818
import numpy as np
1919
from numpy import (
20-
number,
2120
uint64,
2221
int_,
2322
int64,
@@ -48,11 +47,11 @@ from numpy._typing import (
4847
_ShapeLike,
4948
_ArrayLikeBool_co,
5049
_ArrayLikeUInt_co,
50+
_ArrayLikeInt,
5151
_ArrayLikeInt_co,
5252
_ArrayLikeFloat_co,
5353
_ArrayLikeComplex_co,
5454
_ArrayLikeObject_co,
55-
_ArrayLikeTD64_co,
5655
_IntLike_co,
5756
_BoolLike_co,
5857
_ComplexLike_co,
@@ -323,31 +322,42 @@ def matrix_transpose(x: _ArrayLike[_ScalarT], /) -> NDArray[_ScalarT]: ...
323322
@overload
324323
def matrix_transpose(x: ArrayLike, /) -> NDArray[Any]: ...
325324

325+
#
326326
@overload
327327
def partition(
328328
a: _ArrayLike[_ScalarT],
329-
kth: _ArrayLikeInt_co,
330-
axis: SupportsIndex | None = ...,
331-
kind: _PartitionKind = ...,
332-
order: str | Sequence[str] | None = ...,
329+
kth: _ArrayLikeInt,
330+
axis: SupportsIndex | None = -1,
331+
kind: _PartitionKind = "introselect",
332+
order: None = None,
333333
) -> NDArray[_ScalarT]: ...
334334
@overload
335+
def partition(
336+
a: _ArrayLike[np.void],
337+
kth: _ArrayLikeInt,
338+
axis: SupportsIndex | None = -1,
339+
kind: _PartitionKind = "introselect",
340+
order: str | Sequence[str] | None = None,
341+
) -> NDArray[np.void]: ...
342+
@overload
335343
def partition(
336344
a: ArrayLike,
337-
kth: _ArrayLikeInt_co,
338-
axis: SupportsIndex | None = ...,
339-
kind: _PartitionKind = ...,
340-
order: str | Sequence[str] | None = ...,
345+
kth: _ArrayLikeInt,
346+
axis: SupportsIndex | None = -1,
347+
kind: _PartitionKind = "introselect",
348+
order: str | Sequence[str] | None = None,
341349
) -> NDArray[Any]: ...
342350

351+
#
343352
def argpartition(
344353
a: ArrayLike,
345-
kth: _ArrayLikeInt_co,
354+
kth: _ArrayLikeInt,
346355
axis: SupportsIndex | None = -1,
347-
kind: _PartitionKind = ...,
348-
order: str | Sequence[str] | None = ...,
356+
kind: _PartitionKind = "introselect",
357+
order: str | Sequence[str] | None = None,
349358
) -> NDArray[intp]: ...
350359

360+
#
351361
@overload
352362
def sort(
353363
a: _ArrayLike[_ScalarT],

‎numpy/ma/core.pyi

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ from typing import Any, Literal, SupportsIndex, TypeAlias, TypeVar, overload
77
from _typeshed import Incomplete
88
from typing_extensions import deprecated
99

10+
import numpy as np
1011
from numpy import (
1112
_ModeKind,
1213
_OrderKACF,
@@ -688,19 +689,43 @@ class MaskedArray(ndarray[_ShapeT_co, _DTypeT_co]):
688689
) -> _ArrayT: ...
689690

690691
#
692+
@overload
691693
def partition(
692694
self,
695+
/,
693696
kth: _ArrayLikeInt,
694697
axis: SupportsIndex = -1,
695698
kind: _PartitionKind = "introselect",
696-
order: str | Sequence[str] | None = None
699+
order: None = None
697700
) -> None: ...
701+
@overload
702+
def partition(
703+
self: _MaskedArray[np.void],
704+
/,
705+
kth: _ArrayLikeInt,
706+
axis: SupportsIndex = -1,
707+
kind: _PartitionKind = "introselect",
708+
order: str | Sequence[str] | None = None,
709+
) -> None: ...
710+
711+
#
712+
@overload
698713
def argpartition(
699714
self,
715+
/,
700716
kth: _ArrayLikeInt,
701-
axis: SupportsIndex = -1,
717+
axis: SupportsIndex | None = -1,
702718
kind: _PartitionKind = "introselect",
703-
order: str | Sequence[str] | None = None
719+
order: None = None,
720+
) -> _MaskedArray[intp]: ...
721+
@overload
722+
def argpartition(
723+
self: _MaskedArray[np.void],
724+
/,
725+
kth: _ArrayLikeInt,
726+
axis: SupportsIndex | None = -1,
727+
kind: _PartitionKind = "introselect",
728+
order: str | Sequence[str] | None = None,
704729
) -> _MaskedArray[intp]: ...
705730

706731
# Keep in-sync with np.ma.take

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ MAR_td64: MaskedNDArray[np.timedelta64]
2424
MAR_o: MaskedNDArray[np.object_]
2525
MAR_s: MaskedNDArray[np.str_]
2626
MAR_byte: MaskedNDArray[np.bytes_]
27+
MAR_V: MaskedNDArray[np.void]
2728

2829
MAR_subclass: MaskedNDArraySubclass
2930

@@ -163,7 +164,7 @@ assert_type(np.ma.take([1], [0]), MaskedNDArray[Any])
163164
assert_type(np.ma.take(np.eye(2), 1, axis=0), MaskedNDArray[np.float64])
164165

165166
assert_type(MAR_f4.partition(1), None)
166-
assert_type(MAR_f4.partition(1, axis=0, kind='introselect', order='K'), None)
167+
assert_type(MAR_V.partition(1, axis=0, kind='introselect', order='K'), None)
167168

168169
assert_type(MAR_f4.argpartition(1), MaskedNDArray[np.intp])
169170
assert_type(MAR_1d.argpartition(1, axis=0, kind='introselect', order='K'), MaskedNDArray[np.intp])

0 commit comments

Comments
 (0)
0