Skip to content

Commit 58de93c

Browse files
MarcoGorelliMarcoGorelli
authored
Fix erroneous _ListLike type in Series / Index methods (#1426)
* use AnyArrayLike | SequenceNotStr for searchsorted * wip * rename `_ListLike` to `_DataLike`, use `ListLike` instead * extra test# * sorter too * preserve what `_ListLike` used to be for `_DataLike` * fixup * fixup * simplify `ListLike | Index | Series` to `ListLike` * fixup
1 parent 71b5399 commit 58de93c

File tree

4 files changed

+60
-50
lines changed

4 files changed

+60
-50
lines changed

‎pandas-stubs/core/base.pyi‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,16 +28,15 @@ from pandas._libs.tslibs.timedeltas import Timedelta
2828
from pandas._typing import (
2929
S1,
3030
S2,
31-
ArrayLike,
3231
AxisIndex,
3332
DropKeep,
3433
DTypeLike,
3534
GenericT,
3635
GenericT_co,
3736
Just,
37+
ListLike,
3838
NDFrameT,
3939
Scalar,
40-
SequenceNotStr,
4140
SupportsDType,
4241
np_1darray,
4342
np_ndarray_anyint,
@@ -47,8 +46,6 @@ from pandas._typing import (
4746
)
4847
from pandas.util._decorators import cache_readonly
4948

50-
_ListLike: TypeAlias = ArrayLike | dict[str, np.ndarray] | SequenceNotStr[S1]
51-
5249
class NoNewAttributesMixin:
5350
def __setattr__(self, key: str, value: Any) -> None: ...
5451

@@ -157,16 +154,16 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
157154
@overload
158155
def searchsorted(
159156
self,
160-
value: _ListLike,
157+
value: ListLike,
161158
side: Literal["left", "right"] = ...,
162-
sorter: _ListLike | None = ...,
159+
sorter: ListLike | None = None,
163160
) -> np_1darray[np.intp]: ...
164161
@overload
165162
def searchsorted(
166163
self,
167164
value: Scalar,
168165
side: Literal["left", "right"] = ...,
169-
sorter: _ListLike | None = ...,
166+
sorter: ListLike | None = None,
170167
) -> np.intp: ...
171168
def drop_duplicates(self, *, keep: DropKeep = ...) -> Self: ...
172169

‎pandas-stubs/core/indexes/base.pyi‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@ from pandas.core.base import (
4747
Supports_ProtoMul,
4848
Supports_ProtoRAdd,
4949
Supports_ProtoRMul,
50-
_ListLike,
5150
)
5251
from pandas.core.indexes.category import CategoricalIndex
5352
from pandas.core.strings.accessor import StringMethods
@@ -506,7 +505,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
506505
@overload
507506
def __add__(self: Index[Never], other: _str) -> Never: ...
508507
@overload
509-
def __add__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ...
508+
def __add__(
509+
self: Index[Never], other: complex | ArrayLike | SequenceNotStr[S1] | Index
510+
) -> Index: ...
510511
@overload
511512
def __add__(self, other: Index[Never]) -> Index: ...
512513
@overload
@@ -556,7 +557,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
556557
@overload
557558
def __radd__(self: Index[Never], other: _str) -> Never: ...
558559
@overload
559-
def __radd__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ...
560+
def __radd__(
561+
self: Index[Never], other: complex | ArrayLike | SequenceNotStr[S1] | Index
562+
) -> Index: ...
560563
@overload
561564
def __radd__(
562565
self: Supports_ProtoRAdd[_T_contra, S2],
@@ -605,7 +608,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
605608
@overload
606609
def __sub__(self: Index[Never], other: DatetimeIndex) -> Never: ...
607610
@overload
608-
def __sub__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ...
611+
def __sub__(
612+
self: Index[Never], other: complex | ArrayLike | SequenceNotStr[S1] | Index
613+
) -> Index: ...
609614
@overload
610615
def __sub__(self, other: Index[Never]) -> Index: ...
611616
@overload
@@ -741,7 +746,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
741746
),
742747
) -> Index[complex]: ...
743748
@overload
744-
def __mul__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ...
749+
def __mul__(
750+
self: Index[Never], other: complex | ArrayLike | SequenceNotStr[S1] | Index
751+
) -> Index: ...
745752
@overload
746753
def __mul__(self, other: Index[Never]) -> Index: ...
747754
@overload
@@ -817,7 +824,9 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
817824
self: Index[T_COMPLEX], other: np_ndarray_complex | Index[complex]
818825
) -> Index[complex]: ...
819826
@overload
820-
def __rmul__(self: Index[Never], other: complex | _ListLike | Index) -> Index: ...
827+
def __rmul__(
828+
self: Index[Never], other: complex | ArrayLike | SequenceNotStr[S1] | Index
829+
) -> Index: ...
821830
@overload
822831
def __rmul__(self, other: Index[Never]) -> Index: ...
823832
@overload

‎pandas-stubs/core/series.pyi‎

Lines changed: 33 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,6 @@ from pandas.core.base import (
7272
Supports_ProtoMul,
7373
Supports_ProtoRAdd,
7474
Supports_ProtoRMul,
75-
_ListLike,
7675
)
7776
from pandas.core.frame import DataFrame
7877
from pandas.core.generic import NDFrame
@@ -296,7 +295,8 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]):
296295
value: S1 | ArrayLike | Series[S1] | None,
297296
) -> None: ...
298297

299-
_ListLikeS1: TypeAlias = (
298+
_DataLike: TypeAlias = ArrayLike | dict[str, np.ndarray] | SequenceNotStr[S1]
299+
_DataLikeS1: TypeAlias = (
300300
ArrayLike | dict[_str, np.ndarray] | Sequence[S1] | IndexOpsMixin[S1]
301301
)
302302

@@ -360,7 +360,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
360360
@overload
361361
def __new__(
362362
cls,
363-
data: _ListLike,
363+
data: _DataLike,
364364
index: AxesData | None = ...,
365365
*,
366366
dtype: TimestampDtypeArg,
@@ -370,7 +370,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
370370
@overload
371371
def __new__(
372372
cls,
373-
data: _ListLike,
373+
data: _DataLike,
374374
index: AxesData | None = ...,
375375
*,
376376
dtype: CategoryDtypeArg,
@@ -427,7 +427,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
427427
@overload
428428
def __new__( # type: ignore[overload-overlap]
429429
cls,
430-
data: Scalar | _ListLike | dict[HashableT1, Any] | None,
430+
data: Scalar | _DataLike | dict[HashableT1, Any] | None,
431431
index: AxesData | None = ...,
432432
*,
433433
dtype: type[S1],
@@ -474,7 +474,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
474474
def __new__(
475475
cls,
476476
data: (
477-
S1 | _ListLikeS1[S1] | dict[HashableT1, S1] | KeysView[S1] | ValuesView[S1]
477+
S1 | _DataLikeS1[S1] | dict[HashableT1, S1] | KeysView[S1] | ValuesView[S1]
478478
),
479479
index: AxesData | None = ...,
480480
dtype: Dtype = ...,
@@ -486,7 +486,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
486486
cls,
487487
data: (
488488
Scalar
489-
| _ListLike
489+
| _DataLike
490490
| Mapping[HashableT1, Any]
491491
| BaseGroupBy
492492
| NaTType
@@ -853,7 +853,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
853853
@overload
854854
def quantile(
855855
self,
856-
q: _ListLike,
856+
q: ListLike,
857857
interpolation: QuantileInterpolation = ...,
858858
) -> Series[S1]: ...
859859
def corr(
@@ -905,16 +905,16 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
905905
@overload
906906
def searchsorted(
907907
self,
908-
value: _ListLike,
908+
value: ListLike,
909909
side: Literal["left", "right"] = ...,
910-
sorter: _ListLike | None = ...,
910+
sorter: ListLike | None = None,
911911
) -> np_1darray[np.intp]: ...
912912
@overload
913913
def searchsorted(
914914
self,
915915
value: Scalar,
916916
side: Literal["left", "right"] = ...,
917-
sorter: _ListLike | None = ...,
917+
sorter: ListLike | None = None,
918918
) -> np.intp: ...
919919
@overload
920920
def compare(
@@ -1311,7 +1311,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
13111311
) -> Series[S1]: ...
13121312
def filter(
13131313
self,
1314-
items: _ListLike | None = None,
1314+
items: ListLike | None = None,
13151315
like: _str | None = None,
13161316
regex: _str | None = None,
13171317
axis: AxisIndex | None = None,
@@ -1326,7 +1326,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
13261326
n: int | None = None,
13271327
frac: float | None = None,
13281328
replace: _bool = False,
1329-
weights: _str | _ListLike | np.ndarray | None = None,
1329+
weights: _str | ListLike | np.ndarray | None = None,
13301330
random_state: RandomState | None = None,
13311331
axis: AxisIndex | None = None,
13321332
ignore_index: _bool = False,
@@ -1687,9 +1687,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
16871687
# just failed to generate these so I couldn't match
16881688
# them up.
16891689
@overload
1690-
def __add__(
1691-
self: Series[Never], other: complex | _ListLike | Index | Series
1692-
) -> Series: ...
1690+
def __add__(self: Series[Never], other: complex | ListLike) -> Series: ...
16931691
@overload
16941692
def __add__(self, other: Index[Never] | Series[Never]) -> Series: ...
16951693
@overload
@@ -1775,7 +1773,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
17751773
@overload
17761774
def add(
17771775
self: Series[Never],
1778-
other: complex | _ListLike | Index | Series,
1776+
other: complex | ListLike,
17791777
level: Level | None = None,
17801778
fill_value: float | None = None,
17811779
axis: int = 0,
@@ -1906,9 +1904,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
19061904
axis: int = 0,
19071905
) -> Series[_str]: ...
19081906
@overload
1909-
def __radd__(
1910-
self: Series[Never], other: complex | _ListLike | Index | Series
1911-
) -> Series: ...
1907+
def __radd__(self: Series[Never], other: complex | ListLike) -> Series: ...
19121908
@overload
19131909
def __radd__(self, other: Index[Never] | Series[Never]) -> Series: ...
19141910
@overload
@@ -2008,7 +2004,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
20082004
@overload
20092005
def radd(
20102006
self: Series[Never],
2011-
other: complex | _ListLike | Index | Series,
2007+
other: complex | ListLike,
20122008
level: Level | None = None,
20132009
fill_value: float | None = None,
20142010
axis: int = 0,
@@ -2461,16 +2457,16 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
24612457
axis: AxisIndex = ...,
24622458
) -> Series[int]: ...
24632459
def __ge__( # type: ignore[override]
2464-
self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date
2460+
self, other: S1 | ListLike | Series[S1] | datetime | timedelta | date
24652461
) -> Series[_bool]: ...
24662462
def __gt__( # type: ignore[override]
2467-
self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date
2463+
self, other: S1 | ListLike | Series[S1] | datetime | timedelta | date
24682464
) -> Series[_bool]: ...
24692465
def __le__( # type: ignore[override]
2470-
self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date
2466+
self, other: S1 | ListLike | Series[S1] | datetime | timedelta | date
24712467
) -> Series[_bool]: ...
24722468
def __lt__( # type: ignore[override]
2473-
self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date
2469+
self, other: S1 | ListLike | Series[S1] | datetime | timedelta | date
24742470
) -> Series[_bool]: ...
24752471
@overload
24762472
def __mul__( # type: ignore[overload-overlap]
@@ -2567,7 +2563,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
25672563
@overload
25682564
def mul(
25692565
self: Series[Never],
2570-
other: complex | _ListLike | Index | Series,
2566+
other: complex | ListLike,
25712567
level: Level | None = None,
25722568
fill_value: float | None = None,
25732569
axis: int = 0,
@@ -2777,7 +2773,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
27772773
@overload
27782774
def rmul(
27792775
self: Series[Never],
2780-
other: complex | _ListLike | Index | Series,
2776+
other: complex | ListLike,
27812777
level: Level | None = None,
27822778
fill_value: float | None = None,
27832779
axis: int = 0,
@@ -2892,9 +2888,9 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
28922888
fill_value: float | None = None,
28932889
axis: int = 0,
28942890
) -> Series[complex]: ...
2895-
def __mod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
2891+
def __mod__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ...
28962892
def __ne__(self, other: object) -> Series[_bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
2897-
def __pow__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
2893+
def __pow__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ...
28982894
# ignore needed for mypy as we want different results based on the arguments
28992895
@overload # type: ignore[override]
29002896
# pyrefly: ignore # bad-override
@@ -2911,9 +2907,9 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
29112907
) -> Series[bool]: ...
29122908
@overload
29132909
def __rand__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
2914-
def __rdivmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
2915-
def __rmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
2916-
def __rpow__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
2910+
def __rdivmod__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
2911+
def __rmod__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ...
2912+
def __rpow__(self, other: num | ListLike | Series[S1]) -> Series[S1]: ...
29172913
# ignore needed for mypy as we want different results based on the arguments
29182914
@overload # type: ignore[override]
29192915
# pyrefly: ignore # bad-override
@@ -3683,7 +3679,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
36833679
@overload
36843680
def truediv(
36853681
self: Series[Never],
3686-
other: complex | _ListLike | Index | Series,
3682+
other: complex | ListLike,
36873683
level: Level | None = None,
36883684
fill_value: float | None = None,
36893685
axis: AxisIndex = 0,
@@ -3946,7 +3942,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
39463942
@overload
39473943
def rtruediv(
39483944
self: Series[Never],
3949-
other: complex | _ListLike | Index | Series,
3945+
other: complex | ListLike,
39503946
level: Level | None = None,
39513947
fill_value: float | None = None,
39523948
axis: AxisIndex = 0,
@@ -4165,7 +4161,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
41654161
) -> Series[S1]: ...
41664162
def divmod(
41674163
self,
4168-
other: num | _ListLike | Series[S1],
4164+
other: num | ListLike | Series[S1],
41694165
level: Level | None = ...,
41704166
fill_value: float | None = None,
41714167
axis: AxisIndex = ...,
@@ -4323,7 +4319,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
43234319
) -> S1: ...
43244320
def mod(
43254321
self,
4326-
other: num | _ListLike | Series[S1],
4322+
other: num | ListLike | Series[S1],
43274323
level: Level | None = ...,
43284324
fill_value: float | None = None,
43294325
axis: AxisIndex | None = 0,
@@ -4339,7 +4335,7 @@ class Series(IndexOpsMixin[S1], ElementOpsMixin[S1], NDFrame):
43394335
def nunique(self, dropna: _bool = True) -> int: ...
43404336
def pow(
43414337
self,
4342-
other: num | _ListLike | Series[S1],
4338+
other: num | ListLike | Series[S1],
43434339
level: Level | None = ...,
43444340
fill_value: float | None = None,
43454341
axis: AxisIndex | None = 0,

‎tests/indexes/test_indexes.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1500,8 +1500,16 @@ def test_index_searchsorted() -> None:
15001500
idx = pd.Index([1, 2, 3])
15011501
check(assert_type(idx.searchsorted(1), np.intp), np.intp)
15021502
check(assert_type(idx.searchsorted([1]), "np_1darray[np.intp]"), np_1darray)
1503+
check(assert_type(idx.searchsorted(range(1, 2)), "np_1darray[np.intp]"), np_1darray)
1504+
check(
1505+
assert_type(idx.searchsorted(pd.Series([1])), "np_1darray[np.intp]"), np_1darray
1506+
)
1507+
check(
1508+
assert_type(idx.searchsorted(np.array([1])), "np_1darray[np.intp]"), np_1darray
1509+
)
15031510
check(assert_type(idx.searchsorted(1, side="left"), np.intp), np.intp)
15041511
check(assert_type(idx.searchsorted(1, sorter=[1, 0, 2]), np.intp), np.intp)
1512+
check(assert_type(idx.searchsorted(1, sorter=range(3)), np.intp), np.intp)
15051513

15061514

15071515
def test_period_index_constructor() -> None:

0 commit comments

Comments
 (0)