Skip to content

Commit 928df7b

Browse files
cmp0xffcmp0xff
authored
feat(index): ➖ arithmetic subtraction (#1360)
* feat(index): arithmetic subtraction * refactor(series): remove duplicated tests * fix(comment): #1360 * fix(ty): #1360 make it happy
1 parent ffa88e5 commit 928df7b

File tree

16 files changed

+817
-83
lines changed

16 files changed

+817
-83
lines changed

‎pandas-stubs/core/arraylike.pyi

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,6 @@ class OpsMixin:
1919
def __rxor__(self, other: Any) -> Self: ...
2020
# -------------------------------------------------------------
2121
# Arithmetic Methods
22-
def __sub__(self, other: Any) -> Self: ...
23-
def __rsub__(self, other: Any) -> Self: ...
2422
def __mul__(self, other: Any) -> Self: ...
2523
def __rmul__(self, other: Any) -> Self: ...
2624
# Handled by subclasses that specify only the valid values

‎pandas-stubs/core/base.pyi

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from collections.abc import (
22
Hashable,
33
Iterator,
4+
Sequence,
45
)
56
from typing import (
67
Any,
@@ -34,6 +35,10 @@ from pandas._typing import (
3435
SequenceNotStr,
3536
SupportsDType,
3637
np_1darray,
38+
np_ndarray_anyint,
39+
np_ndarray_bool,
40+
np_ndarray_complex,
41+
np_ndarray_float,
3742
)
3843
from pandas.util._decorators import cache_readonly
3944

@@ -160,3 +165,14 @@ class IndexOpsMixin(OpsMixin, Generic[S1, GenericT_co]):
160165
sorter: _ListLike | None = ...,
161166
) -> np.intp: ...
162167
def drop_duplicates(self, *, keep: DropKeep = ...) -> Self: ...
168+
169+
NumListLike: TypeAlias = (
170+
ExtensionArray
171+
| np_ndarray_bool
172+
| np_ndarray_anyint
173+
| np_ndarray_float
174+
| np_ndarray_complex
175+
| dict[str, np.ndarray]
176+
| Sequence[complex]
177+
| IndexOpsMixin[complex]
178+
)

‎pandas-stubs/core/frame.pyi

Lines changed: 16 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1788,6 +1788,22 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
17881788
level: Level | None = None,
17891789
fill_value: float | None = None,
17901790
) -> Self: ...
1791+
def __sub__(self, other: Any) -> Self: ...
1792+
def sub(
1793+
self,
1794+
other: num | ListLike | DataFrame,
1795+
axis: Axis | None = ...,
1796+
level: Level | None = ...,
1797+
fill_value: float | None = None,
1798+
) -> Self: ...
1799+
def __rsub__(self, other: Any) -> Self: ...
1800+
def rsub(
1801+
self,
1802+
other,
1803+
axis: Axis = ...,
1804+
level: Level | None = ...,
1805+
fill_value: float | None = None,
1806+
) -> Self: ...
17911807
@final
17921808
def add_prefix(self, prefix: _str, axis: Axis | None = None) -> Self: ...
17931809
@final
@@ -2353,13 +2369,6 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
23532369
level: Level | None = ...,
23542370
fill_value: float | None = None,
23552371
) -> Self: ...
2356-
def rsub(
2357-
self,
2358-
other,
2359-
axis: Axis = ...,
2360-
level: Level | None = ...,
2361-
fill_value: float | None = None,
2362-
) -> Self: ...
23632372
def rtruediv(
23642373
self,
23652374
other,
@@ -2405,20 +2414,6 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
24052414
numeric_only: _bool = False,
24062415
**kwargs: Any,
24072416
) -> Series: ...
2408-
def sub(
2409-
self,
2410-
other: num | ListLike | DataFrame,
2411-
axis: Axis | None = ...,
2412-
level: Level | None = ...,
2413-
fill_value: float | None = None,
2414-
) -> Self: ...
2415-
def subtract(
2416-
self,
2417-
other: num | ListLike | DataFrame,
2418-
axis: Axis | None = ...,
2419-
level: Level | None = ...,
2420-
fill_value: float | None = None,
2421-
) -> Self: ...
24222417
def sum(
24232418
self,
24242419
axis: Axis = 0,

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

Lines changed: 144 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,12 +34,15 @@ from pandas import (
3434
TimedeltaIndex,
3535
)
3636
from pandas.core.arrays import ExtensionArray
37-
from pandas.core.base import IndexOpsMixin
37+
from pandas.core.base import (
38+
IndexOpsMixin,
39+
NumListLike,
40+
_ListLike,
41+
)
3842
from pandas.core.strings.accessor import StringMethods
3943
from typing_extensions import (
4044
Never,
4145
Self,
42-
TypeAlias,
4346
)
4447

4548
from pandas._libs.interval import _OrderableT
@@ -60,6 +63,7 @@ from pandas._typing import (
6063
GenericT_co,
6164
HashableT,
6265
IgnoreRaise,
66+
Just,
6367
Label,
6468
Level,
6569
MaskType,
@@ -82,8 +86,6 @@ from pandas._typing import (
8286

8387
class InvalidIndexError(Exception): ...
8488

85-
_ListLike: TypeAlias = ArrayLike | dict[_str, np.ndarray] | SequenceNotStr[S1]
86-
8789
class Index(IndexOpsMixin[S1]):
8890
__hash__: ClassVar[None] # type: ignore[assignment]
8991
# overloads with additional dtypes
@@ -626,6 +628,144 @@ class Index(IndexOpsMixin[S1]):
626628
self: Index[_str], other: _str | Sequence[_str] | np_ndarray_str | Index[_str]
627629
) -> Index[_str]: ...
628630
@overload
631+
def __sub__(self: Index[Never], other: DatetimeIndex) -> Never: ...
632+
@overload
633+
def __sub__(self: Index[Never], other: complex | NumListLike | Index) -> Index: ...
634+
@overload
635+
def __sub__(self, other: Index[Never]) -> Index: ...
636+
@overload
637+
def __sub__(
638+
self: Index[bool],
639+
other: Just[int] | Sequence[Just[int]] | np_ndarray_anyint | Index[int],
640+
) -> Index[int]: ...
641+
@overload
642+
def __sub__(
643+
self: Index[bool],
644+
other: Just[float] | Sequence[Just[float]] | np_ndarray_float | Index[float],
645+
) -> Index[float]: ...
646+
@overload
647+
def __sub__(
648+
self: Index[int],
649+
other: (
650+
int
651+
| Sequence[int]
652+
| np_ndarray_bool
653+
| np_ndarray_anyint
654+
| Index[bool]
655+
| Index[int]
656+
),
657+
) -> Index[int]: ...
658+
@overload
659+
def __sub__(
660+
self: Index[int],
661+
other: Just[float] | Sequence[Just[float]] | np_ndarray_float | Index[float],
662+
) -> Index[float]: ...
663+
@overload
664+
def __sub__(
665+
self: Index[float],
666+
other: (
667+
float
668+
| Sequence[float]
669+
| np_ndarray_bool
670+
| np_ndarray_anyint
671+
| np_ndarray_float
672+
| Index[bool]
673+
| Index[int]
674+
| Index[float]
675+
),
676+
) -> Index[float]: ...
677+
@overload
678+
def __sub__(
679+
self: Index[complex],
680+
other: (
681+
T_COMPLEX
682+
| Sequence[T_COMPLEX]
683+
| np_ndarray_bool
684+
| np_ndarray_anyint
685+
| np_ndarray_float
686+
| Index[T_COMPLEX]
687+
),
688+
) -> Index[complex]: ...
689+
@overload
690+
def __sub__(
691+
self: Index[T_COMPLEX],
692+
other: (
693+
Just[complex]
694+
| Sequence[Just[complex]]
695+
| np_ndarray_complex
696+
| Index[complex]
697+
),
698+
) -> Index[complex]: ...
699+
@overload
700+
def __rsub__(self: Index[Never], other: DatetimeIndex) -> Never: ... # type: ignore[misc]
701+
@overload
702+
def __rsub__(self: Index[Never], other: complex | NumListLike | Index) -> Index: ...
703+
@overload
704+
def __rsub__(self, other: Index[Never]) -> Index: ...
705+
@overload
706+
def __rsub__(
707+
self: Index[bool],
708+
other: Just[int] | Sequence[Just[int]] | np_ndarray_anyint | Index[int],
709+
) -> Index[int]: ...
710+
@overload
711+
def __rsub__(
712+
self: Index[bool],
713+
other: Just[float] | Sequence[Just[float]] | np_ndarray_float | Index[float],
714+
) -> Index[float]: ...
715+
@overload
716+
def __rsub__(
717+
self: Index[int],
718+
other: (
719+
int
720+
| Sequence[int]
721+
| np_ndarray_bool
722+
| np_ndarray_anyint
723+
| Index[bool]
724+
| Index[int]
725+
),
726+
) -> Index[int]: ...
727+
@overload
728+
def __rsub__(
729+
self: Index[int],
730+
other: Just[float] | Sequence[Just[float]] | np_ndarray_float | Index[float],
731+
) -> Index[float]: ...
732+
@overload
733+
def __rsub__(
734+
self: Index[float],
735+
other: (
736+
float
737+
| Sequence[float]
738+
| np_ndarray_bool
739+
| np_ndarray_anyint
740+
| np_ndarray_float
741+
| Index[bool]
742+
| Index[int]
743+
| Index[float]
744+
),
745+
) -> Index[float]: ...
746+
@overload
747+
def __rsub__(
748+
self: Index[complex],
749+
other: (
750+
T_COMPLEX
751+
| Sequence[T_COMPLEX]
752+
| np_ndarray_bool
753+
| np_ndarray_anyint
754+
| np_ndarray_float
755+
| Index[T_COMPLEX]
756+
),
757+
) -> Index[complex]: ...
758+
@overload
759+
def __rsub__(
760+
self: Index[T_COMPLEX],
761+
other: (
762+
Just[complex]
763+
| Sequence[Just[complex]]
764+
| np_ndarray_complex
765+
| Index[complex]
766+
),
767+
) -> Index[complex]: ...
768+
@overload
629769
def __mul__(
630770
self: Index[int] | Index[float], other: timedelta
631771
) -> TimedeltaIndex: ...

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ class DatetimeIndexOpsMixin(ExtensionIndex[S1, GenericT_co]):
3030
def argmax(
3131
self, axis: AxisIndex | None = None, skipna: bool = True, *args, **kwargs
3232
) -> np.int64: ...
33-
def __rsub__( # type: ignore[override]
33+
def __rsub__( # type: ignore[misc,override] # pyright: ignore[reportIncompatibleMethodOverride]
3434
self, other: DatetimeIndexOpsMixin
3535
) -> TimedeltaIndex: ...
3636

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

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ from pandas._typing import (
3737
IntervalClosedType,
3838
TimeUnit,
3939
TimeZones,
40+
np_ndarray_dt,
41+
np_ndarray_td,
4042
)
4143

4244
from pandas.core.dtypes.dtypes import DatetimeTZDtype
@@ -67,15 +69,16 @@ class DatetimeIndex(
6769
def __add__( # pyright: ignore[reportIncompatibleMethodOverride]
6870
self, other: timedelta | Timedelta | TimedeltaIndex | BaseOffset
6971
) -> DatetimeIndex: ...
70-
@overload
72+
@overload # type: ignore[override]
7173
def __sub__(self, other: TimedeltaSeries) -> TimestampSeries: ...
7274
@overload
7375
def __sub__(
74-
self, other: timedelta | Timedelta | TimedeltaIndex | BaseOffset
76+
self,
77+
other: timedelta | np.timedelta64 | np_ndarray_td | TimedeltaIndex | BaseOffset,
7578
) -> DatetimeIndex: ...
7679
@overload
77-
def __sub__(
78-
self, other: datetime | Timestamp | DatetimeIndex
80+
def __sub__( # pyright: ignore[reportIncompatibleMethodOverride]
81+
self, other: datetime | np.datetime64 | np_ndarray_dt | DatetimeIndex
7982
) -> TimedeltaIndex: ...
8083
@final
8184
def to_series(self, index=..., name: Hashable = ...) -> TimestampSeries: ...

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexField
3636
) -> Self: ...
3737
@property
3838
def values(self) -> np_1darray[np.object_]: ...
39-
@overload
39+
@overload # type: ignore[override]
4040
def __sub__(self, other: Period) -> Index: ...
4141
@overload
4242
def __sub__(self, other: Self) -> Index: ...
@@ -45,7 +45,9 @@ class PeriodIndex(DatetimeIndexOpsMixin[pd.Period, np.object_], PeriodIndexField
4545
@overload
4646
def __sub__(self, other: NaTType) -> NaTType: ...
4747
@overload
48-
def __sub__(self, other: TimedeltaIndex | pd.Timedelta) -> Self: ...
48+
def __sub__( # pyright: ignore[reportIncompatibleMethodOverride]
49+
self, other: TimedeltaIndex | pd.Timedelta
50+
) -> Self: ...
4951
@overload # type: ignore[override]
5052
def __rsub__(self, other: Period) -> Index: ...
5153
@overload

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ from pandas._libs.tslibs import BaseOffset
3030
from pandas._typing import (
3131
AxesData,
3232
TimedeltaConvertibleTypes,
33+
np_ndarray_td,
3334
num,
3435
)
3536

@@ -58,7 +59,9 @@ class TimedeltaIndex(
5859
self, other: dt.timedelta | Timedelta | Self
5960
) -> Self: ...
6061
def __radd__(self, other: dt.datetime | Timestamp | DatetimeIndex) -> DatetimeIndex: ... # type: ignore[override]
61-
def __sub__(self, other: dt.timedelta | Timedelta | Self) -> Self: ...
62+
def __sub__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
63+
self, other: dt.timedelta | np.timedelta64 | np_ndarray_td | Self
64+
) -> Self: ...
6265
def __mul__(self, other: num) -> Self: ...
6366
@overload # type: ignore[override]
6467
def __truediv__(self, other: num | Sequence[float]) -> Self: ...

0 commit comments

Comments
 (0)