Skip to content

Commit 1400649

Browse files
authored
TYP: sort_index (#46300)
1 parent 2cb628d commit 1400649

File tree

5 files changed

+207
-44
lines changed

5 files changed

+207
-44
lines changed

pandas/_typing.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -307,10 +307,12 @@ def closed(self) -> bool:
307307
XMLParsers = Literal["lxml", "etree"]
308308

309309
# Interval closed type
310-
311310
IntervalClosedType = Literal["left", "right", "both", "neither"]
312311

313312
# datetime and NaTType
314-
315313
DatetimeNaTType = Union[datetime, "NaTType"]
316314
DateTimeErrorChoices = Literal["ignore", "raise", "coerce"]
315+
316+
# sort_index
317+
SortKind = Literal["quicksort", "mergesort", "heapsort", "stable"]
318+
NaPosition = Literal["first", "last"]

pandas/core/frame.py

+65-14
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,12 @@
6161
IndexKeyFunc,
6262
IndexLabel,
6363
Level,
64+
NaPosition,
6465
PythonFuncType,
6566
ReadBuffer,
6667
Renamer,
6768
Scalar,
69+
SortKind,
6870
StorageOptions,
6971
Suffixes,
7072
TimedeltaConvertibleTypes,
@@ -6433,19 +6435,68 @@ def sort_values( # type: ignore[override]
64336435
else:
64346436
return result.__finalize__(self, method="sort_values")
64356437

6436-
@deprecate_nonkeyword_arguments(version=None, allowed_args=["self"])
6438+
@overload
6439+
def sort_index(
6440+
self,
6441+
*,
6442+
axis: Axis = ...,
6443+
level: Level | None = ...,
6444+
ascending: bool | Sequence[bool] = ...,
6445+
inplace: Literal[True],
6446+
kind: SortKind = ...,
6447+
na_position: NaPosition = ...,
6448+
sort_remaining: bool = ...,
6449+
ignore_index: bool = ...,
6450+
key: IndexKeyFunc = ...,
6451+
) -> None:
6452+
...
6453+
6454+
@overload
6455+
def sort_index(
6456+
self,
6457+
*,
6458+
axis: Axis = ...,
6459+
level: Level | None = ...,
6460+
ascending: bool | Sequence[bool] = ...,
6461+
inplace: Literal[False] = ...,
6462+
kind: SortKind = ...,
6463+
na_position: NaPosition = ...,
6464+
sort_remaining: bool = ...,
6465+
ignore_index: bool = ...,
6466+
key: IndexKeyFunc = ...,
6467+
) -> DataFrame:
6468+
...
6469+
6470+
@overload
64376471
def sort_index(
6472+
self,
6473+
*,
6474+
axis: Axis = ...,
6475+
level: Level | None = ...,
6476+
ascending: bool | Sequence[bool] = ...,
6477+
inplace: bool = ...,
6478+
kind: SortKind = ...,
6479+
na_position: NaPosition = ...,
6480+
sort_remaining: bool = ...,
6481+
ignore_index: bool = ...,
6482+
key: IndexKeyFunc = ...,
6483+
) -> DataFrame | None:
6484+
...
6485+
6486+
# error: Signature of "sort_index" incompatible with supertype "NDFrame"
6487+
@deprecate_nonkeyword_arguments(version=None, allowed_args=["self"])
6488+
def sort_index( # type: ignore[override]
64386489
self,
64396490
axis: Axis = 0,
64406491
level: Level | None = None,
6441-
ascending: bool | int | Sequence[bool | int] = True,
6492+
ascending: bool | Sequence[bool] = True,
64426493
inplace: bool = False,
6443-
kind: str = "quicksort",
6444-
na_position: str = "last",
6494+
kind: SortKind = "quicksort",
6495+
na_position: NaPosition = "last",
64456496
sort_remaining: bool = True,
64466497
ignore_index: bool = False,
64476498
key: IndexKeyFunc = None,
6448-
):
6499+
) -> DataFrame | None:
64496500
"""
64506501
Sort object by labels (along an axis).
64516502
@@ -6536,15 +6587,15 @@ def sort_index(
65366587
d 4
65376588
"""
65386589
return super().sort_index(
6539-
axis,
6540-
level,
6541-
ascending,
6542-
inplace,
6543-
kind,
6544-
na_position,
6545-
sort_remaining,
6546-
ignore_index,
6547-
key,
6590+
axis=axis,
6591+
level=level,
6592+
ascending=ascending,
6593+
inplace=inplace,
6594+
kind=kind,
6595+
na_position=na_position,
6596+
sort_remaining=sort_remaining,
6597+
ignore_index=ignore_index,
6598+
key=key,
65486599
)
65496600

65506601
def value_counts(

pandas/core/generic.py

+57-7
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,11 @@
5050
JSONSerializable,
5151
Level,
5252
Manager,
53+
NaPosition,
5354
NDFrameT,
5455
RandomState,
5556
Renamer,
57+
SortKind,
5658
StorageOptions,
5759
T,
5860
TimedeltaConvertibleTypes,
@@ -4683,18 +4685,66 @@ def sort_values(
46834685
"""
46844686
raise AbstractMethodError(self)
46854687

4688+
@overload
46864689
def sort_index(
46874690
self,
4688-
axis=0,
4689-
level=None,
4690-
ascending: bool_t | int | Sequence[bool_t | int] = True,
4691+
*,
4692+
axis: Axis = ...,
4693+
level: Level | None = ...,
4694+
ascending: bool_t | Sequence[bool_t] = ...,
4695+
inplace: Literal[True],
4696+
kind: SortKind = ...,
4697+
na_position: NaPosition = ...,
4698+
sort_remaining: bool_t = ...,
4699+
ignore_index: bool_t = ...,
4700+
key: IndexKeyFunc = ...,
4701+
) -> None:
4702+
...
4703+
4704+
@overload
4705+
def sort_index(
4706+
self: NDFrameT,
4707+
*,
4708+
axis: Axis = ...,
4709+
level: Level | None = ...,
4710+
ascending: bool_t | Sequence[bool_t] = ...,
4711+
inplace: Literal[False] = ...,
4712+
kind: SortKind = ...,
4713+
na_position: NaPosition = ...,
4714+
sort_remaining: bool_t = ...,
4715+
ignore_index: bool_t = ...,
4716+
key: IndexKeyFunc = ...,
4717+
) -> NDFrameT:
4718+
...
4719+
4720+
@overload
4721+
def sort_index(
4722+
self: NDFrameT,
4723+
*,
4724+
axis: Axis = ...,
4725+
level: Level | None = ...,
4726+
ascending: bool_t | Sequence[bool_t] = ...,
4727+
inplace: bool_t = ...,
4728+
kind: SortKind = ...,
4729+
na_position: NaPosition = ...,
4730+
sort_remaining: bool_t = ...,
4731+
ignore_index: bool_t = ...,
4732+
key: IndexKeyFunc = ...,
4733+
) -> NDFrameT | None:
4734+
...
4735+
4736+
def sort_index(
4737+
self: NDFrameT,
4738+
axis: Axis = 0,
4739+
level: Level | None = None,
4740+
ascending: bool_t | Sequence[bool_t] = True,
46914741
inplace: bool_t = False,
4692-
kind: str = "quicksort",
4693-
na_position: str = "last",
4742+
kind: SortKind = "quicksort",
4743+
na_position: NaPosition = "last",
46944744
sort_remaining: bool_t = True,
46954745
ignore_index: bool_t = False,
46964746
key: IndexKeyFunc = None,
4697-
):
4747+
) -> NDFrameT | None:
46984748

46994749
inplace = validate_bool_kwarg(inplace, "inplace")
47004750
axis = self._get_axis_number(axis)
@@ -4715,7 +4765,7 @@ def sort_index(
47154765
if ignore_index:
47164766
result.index = default_index(len(self))
47174767
if inplace:
4718-
return
4768+
return None
47194769
else:
47204770
return result
47214771

pandas/core/series.py

+68-16
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,10 @@
3939
DtypeObj,
4040
FillnaOptions,
4141
IndexKeyFunc,
42+
Level,
43+
NaPosition,
4244
SingleManager,
45+
SortKind,
4346
StorageOptions,
4447
TimedeltaConvertibleTypes,
4548
TimestampConvertibleTypes,
@@ -3575,19 +3578,68 @@ def sort_values(
35753578
else:
35763579
return result.__finalize__(self, method="sort_values")
35773580

3578-
@deprecate_nonkeyword_arguments(version=None, allowed_args=["self"])
3581+
@overload
35793582
def sort_index(
35803583
self,
3581-
axis=0,
3582-
level=None,
3583-
ascending: bool | int | Sequence[bool | int] = True,
3584+
*,
3585+
axis: Axis = ...,
3586+
level: Level | None = ...,
3587+
ascending: bool | Sequence[bool] = ...,
3588+
inplace: Literal[True],
3589+
kind: SortKind = ...,
3590+
na_position: NaPosition = ...,
3591+
sort_remaining: bool = ...,
3592+
ignore_index: bool = ...,
3593+
key: IndexKeyFunc = ...,
3594+
) -> None:
3595+
...
3596+
3597+
@overload
3598+
def sort_index(
3599+
self,
3600+
*,
3601+
axis: Axis = ...,
3602+
level: Level | None = ...,
3603+
ascending: bool | Sequence[bool] = ...,
3604+
inplace: Literal[False] = ...,
3605+
kind: SortKind = ...,
3606+
na_position: NaPosition = ...,
3607+
sort_remaining: bool = ...,
3608+
ignore_index: bool = ...,
3609+
key: IndexKeyFunc = ...,
3610+
) -> Series:
3611+
...
3612+
3613+
@overload
3614+
def sort_index(
3615+
self,
3616+
*,
3617+
axis: Axis = ...,
3618+
level: Level | None = ...,
3619+
ascending: bool | Sequence[bool] = ...,
3620+
inplace: bool = ...,
3621+
kind: SortKind = ...,
3622+
na_position: NaPosition = ...,
3623+
sort_remaining: bool = ...,
3624+
ignore_index: bool = ...,
3625+
key: IndexKeyFunc = ...,
3626+
) -> Series | None:
3627+
...
3628+
3629+
# error: Signature of "sort_index" incompatible with supertype "NDFrame"
3630+
@deprecate_nonkeyword_arguments(version=None, allowed_args=["self"])
3631+
def sort_index( # type: ignore[override]
3632+
self,
3633+
axis: Axis = 0,
3634+
level: Level | None = None,
3635+
ascending: bool | Sequence[bool] = True,
35843636
inplace: bool = False,
3585-
kind: str = "quicksort",
3586-
na_position: str = "last",
3637+
kind: SortKind = "quicksort",
3638+
na_position: NaPosition = "last",
35873639
sort_remaining: bool = True,
35883640
ignore_index: bool = False,
35893641
key: IndexKeyFunc = None,
3590-
):
3642+
) -> Series | None:
35913643
"""
35923644
Sort Series by index labels.
35933645
@@ -3724,15 +3776,15 @@ def sort_index(
37243776
"""
37253777

37263778
return super().sort_index(
3727-
axis,
3728-
level,
3729-
ascending,
3730-
inplace,
3731-
kind,
3732-
na_position,
3733-
sort_remaining,
3734-
ignore_index,
3735-
key,
3779+
axis=axis,
3780+
level=level,
3781+
ascending=ascending,
3782+
inplace=inplace,
3783+
kind=kind,
3784+
na_position=na_position,
3785+
sort_remaining=sort_remaining,
3786+
ignore_index=ignore_index,
3787+
key=key,
37363788
)
37373789

37383790
def argsort(self, axis=0, kind="quicksort", order=None) -> Series:

pandas/core/sorting.py

+13-5
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
Hashable,
1010
Iterable,
1111
Sequence,
12+
cast,
1213
)
1314
import warnings
1415

@@ -22,7 +23,10 @@
2223
from pandas._libs.hashtable import unique_label_indices
2324
from pandas._typing import (
2425
IndexKeyFunc,
26+
Level,
27+
NaPosition,
2528
Shape,
29+
SortKind,
2630
npt,
2731
)
2832

@@ -47,10 +51,10 @@
4751

4852
def get_indexer_indexer(
4953
target: Index,
50-
level: str | int | list[str] | list[int],
51-
ascending: Sequence[bool | int] | bool | int,
52-
kind: str,
53-
na_position: str,
54+
level: Level | list[Level] | None,
55+
ascending: Sequence[bool] | bool,
56+
kind: SortKind,
57+
na_position: NaPosition,
5458
sort_remaining: bool,
5559
key: IndexKeyFunc,
5660
) -> npt.NDArray[np.intp] | None:
@@ -92,8 +96,12 @@ def get_indexer_indexer(
9296
):
9397
return None
9498

99+
# ascending can only be a Sequence for MultiIndex
95100
indexer = nargsort(
96-
target, kind=kind, ascending=ascending, na_position=na_position
101+
target,
102+
kind=kind,
103+
ascending=cast(bool, ascending),
104+
na_position=na_position,
97105
)
98106
return indexer
99107

0 commit comments

Comments
 (0)