Skip to content

Commit 61de1a1

Browse files
cmp0xffcmp0xff
committed
fix: comments
1 parent 44d0be0 commit 61de1a1

File tree

3 files changed

+174
-79
lines changed

3 files changed

+174
-79
lines changed

‎pandas-stubs/_typing.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -820,6 +820,8 @@ np_ndarray_float: TypeAlias = npt.NDArray[np.floating]
820820
np_ndarray_complex: TypeAlias = npt.NDArray[np.complexfloating]
821821
np_ndarray_bool: TypeAlias = npt.NDArray[np.bool_]
822822
np_ndarray_str: TypeAlias = npt.NDArray[np.str_]
823+
np_ndarray_dt: TypeAlias = npt.NDArray[np.datetime64]
824+
np_ndarray_td: TypeAlias = npt.NDArray[np.timedelta64]
823825

824826
# Define shape and generic type variables with defaults similar to numpy
825827
GenericT = TypeVar("GenericT", bound=np.generic, default=Any)

‎pandas-stubs/core/series.pyi

Lines changed: 37 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,7 @@ from pandas._typing import (
186186
np_ndarray_anyint,
187187
np_ndarray_bool,
188188
np_ndarray_complex,
189+
np_ndarray_dt,
189190
np_ndarray_float,
190191
npt,
191192
num,
@@ -2071,6 +2072,11 @@ class Series(IndexOpsMixin[S1], NDFrame):
20712072
@overload
20722073
def __rxor__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
20732074
@overload
2075+
def __sub__(
2076+
self: Series[Never],
2077+
other: datetime | np.datetime64 | np_ndarray_dt | TimestampSeries,
2078+
) -> TimedeltaSeries: ...
2079+
@overload
20742080
def __sub__(self: Series[Never], other: complex | _ListLike | Series) -> Series: ...
20752081
@overload
20762082
def __sub__(self, other: Series[Never]) -> Series: ... # type: ignore[overload-overlap]
@@ -2148,6 +2154,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
21482154
other: timedelta | np.timedelta64 | TimedeltaSeries | TimedeltaIndex,
21492155
) -> TimedeltaSeries: ...
21502156
@overload
2157+
def sub(
2158+
self: Series[Never],
2159+
other: datetime | np.datetime64 | np_ndarray_dt | TimestampSeries,
2160+
level: Level | None = None,
2161+
fill_value: float | None = None,
2162+
axis: int = 0,
2163+
) -> TimedeltaSeries: ...
2164+
@overload
21512165
def sub(
21522166
self: Series[Never],
21532167
other: complex | _ListLike | Series,
@@ -2236,34 +2250,39 @@ class Series(IndexOpsMixin[S1], NDFrame):
22362250
) -> Series[complex]: ...
22372251
@overload
22382252
def sub(
2239-
self: Series[Timestamp],
2240-
other: timedelta | np.timedelta64 | TimedeltaSeries | TimedeltaIndex,
2253+
self: Series[_T_COMPLEX],
2254+
other: (
2255+
Just[complex]
2256+
| Sequence[Just[complex]]
2257+
| np_ndarray_complex
2258+
| Series[complex]
2259+
),
22412260
level: Level | None = None,
22422261
fill_value: float | None = None,
22432262
axis: int = 0,
22442263
) -> Series[complex]: ...
22452264
@overload
22462265
def sub(
2247-
self: Series[Timedelta],
2266+
self: Series[Timestamp],
22482267
other: timedelta | np.timedelta64 | TimedeltaSeries | TimedeltaIndex,
22492268
level: Level | None = None,
22502269
fill_value: float | None = None,
22512270
axis: int = 0,
22522271
) -> Series[complex]: ...
22532272
@overload
22542273
def sub(
2255-
self: Series[_T_COMPLEX],
2256-
other: (
2257-
Just[complex]
2258-
| Sequence[Just[complex]]
2259-
| np_ndarray_complex
2260-
| Series[complex]
2261-
),
2274+
self: Series[Timedelta],
2275+
other: timedelta | np.timedelta64 | TimedeltaSeries | TimedeltaIndex,
22622276
level: Level | None = None,
22632277
fill_value: float | None = None,
22642278
axis: int = 0,
22652279
) -> Series[complex]: ...
22662280
@overload
2281+
def __rsub__( # type: ignore[misc]
2282+
self: Series[Never],
2283+
other: datetime | np.datetime64 | np_ndarray_dt | TimestampSeries,
2284+
) -> TimedeltaSeries: ...
2285+
@overload
22672286
def __rsub__(
22682287
self: Series[Never], other: complex | _ListLike | Series
22692288
) -> Series: ...
@@ -2333,6 +2352,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
23332352
),
23342353
) -> Series[complex]: ...
23352354
@overload
2355+
def rsub(
2356+
self: Series[Never],
2357+
other: datetime | np.datetime64 | np_ndarray_dt | TimestampSeries,
2358+
level: Level | None = None,
2359+
fill_value: float | None = None,
2360+
axis: int = 0,
2361+
) -> TimedeltaSeries: ...
2362+
@overload
23362363
def rsub(
23372364
self: Series[Never],
23382365
other: complex | _ListLike | Series,

‎tests/series/arithmetic/test_sub.py

Lines changed: 135 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,11 @@
1-
from typing import NoReturn
1+
from datetime import (
2+
datetime,
3+
timedelta,
4+
)
5+
from typing import (
6+
TYPE_CHECKING,
7+
NoReturn,
8+
)
29

310
import numpy as np
411
from numpy import typing as npt # noqa: F401
@@ -7,57 +14,60 @@
714

815
from tests import check
916

10-
left = pd.DataFrame({"a": [1, 2, 3]})["a"] # left operand
17+
if TYPE_CHECKING:
18+
from pandas.core.series import TimedeltaSeries # noqa: F401
19+
20+
left_i = pd.DataFrame({"a": [1, 2, 3]})["a"] # left operand
1121

1222

1323
def test_sub_py_scalar() -> None:
1424
"""Test pd.Series[Any] - Python native scalars"""
1525
b, i, f, c = True, 1, 1.0, 1j
1626

17-
check(assert_type(left - b, pd.Series), pd.Series)
18-
check(assert_type(left - i, pd.Series), pd.Series)
19-
check(assert_type(left - f, pd.Series), pd.Series)
20-
check(assert_type(left - c, pd.Series), pd.Series)
27+
check(assert_type(left_i - b, pd.Series), pd.Series)
28+
check(assert_type(left_i - i, pd.Series), pd.Series)
29+
check(assert_type(left_i - f, pd.Series), pd.Series)
30+
check(assert_type(left_i - c, pd.Series), pd.Series)
2131

22-
check(assert_type(b - left, pd.Series), pd.Series)
23-
check(assert_type(i - left, pd.Series), pd.Series)
24-
check(assert_type(f - left, pd.Series), pd.Series)
25-
check(assert_type(c - left, pd.Series), pd.Series)
32+
check(assert_type(b - left_i, pd.Series), pd.Series)
33+
check(assert_type(i - left_i, pd.Series), pd.Series)
34+
check(assert_type(f - left_i, pd.Series), pd.Series)
35+
check(assert_type(c - left_i, pd.Series), pd.Series)
2636

27-
check(assert_type(left.sub(b), pd.Series), pd.Series)
28-
check(assert_type(left.sub(i), pd.Series), pd.Series)
29-
check(assert_type(left.sub(f), pd.Series), pd.Series)
30-
check(assert_type(left.sub(c), pd.Series), pd.Series)
37+
check(assert_type(left_i.sub(b), pd.Series), pd.Series)
38+
check(assert_type(left_i.sub(i), pd.Series), pd.Series)
39+
check(assert_type(left_i.sub(f), pd.Series), pd.Series)
40+
check(assert_type(left_i.sub(c), pd.Series), pd.Series)
3141

32-
check(assert_type(left.rsub(b), pd.Series), pd.Series)
33-
check(assert_type(left.rsub(i), pd.Series), pd.Series)
34-
check(assert_type(left.rsub(f), pd.Series), pd.Series)
35-
check(assert_type(left.rsub(c), pd.Series), pd.Series)
42+
check(assert_type(left_i.rsub(b), pd.Series), pd.Series)
43+
check(assert_type(left_i.rsub(i), pd.Series), pd.Series)
44+
check(assert_type(left_i.rsub(f), pd.Series), pd.Series)
45+
check(assert_type(left_i.rsub(c), pd.Series), pd.Series)
3646

3747

3848
def test_sub_py_sequence() -> None:
3949
"""Test pd.Series[Any] - Python native sequence"""
4050
b, i, f, c = [True, False, True], [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
4151

42-
check(assert_type(left - b, pd.Series), pd.Series)
43-
check(assert_type(left - i, pd.Series), pd.Series)
44-
check(assert_type(left - f, pd.Series), pd.Series)
45-
check(assert_type(left - c, pd.Series), pd.Series)
52+
check(assert_type(left_i - b, pd.Series), pd.Series)
53+
check(assert_type(left_i - i, pd.Series), pd.Series)
54+
check(assert_type(left_i - f, pd.Series), pd.Series)
55+
check(assert_type(left_i - c, pd.Series), pd.Series)
4656

47-
check(assert_type(b - left, pd.Series), pd.Series)
48-
check(assert_type(i - left, pd.Series), pd.Series)
49-
check(assert_type(f - left, pd.Series), pd.Series)
50-
check(assert_type(c - left, pd.Series), pd.Series)
57+
check(assert_type(b - left_i, pd.Series), pd.Series)
58+
check(assert_type(i - left_i, pd.Series), pd.Series)
59+
check(assert_type(f - left_i, pd.Series), pd.Series)
60+
check(assert_type(c - left_i, pd.Series), pd.Series)
5161

52-
check(assert_type(left.sub(b), pd.Series), pd.Series)
53-
check(assert_type(left.sub(i), pd.Series), pd.Series)
54-
check(assert_type(left.sub(f), pd.Series), pd.Series)
55-
check(assert_type(left.sub(c), pd.Series), pd.Series)
62+
check(assert_type(left_i.sub(b), pd.Series), pd.Series)
63+
check(assert_type(left_i.sub(i), pd.Series), pd.Series)
64+
check(assert_type(left_i.sub(f), pd.Series), pd.Series)
65+
check(assert_type(left_i.sub(c), pd.Series), pd.Series)
5666

57-
check(assert_type(left.rsub(b), pd.Series), pd.Series)
58-
check(assert_type(left.rsub(i), pd.Series), pd.Series)
59-
check(assert_type(left.rsub(f), pd.Series), pd.Series)
60-
check(assert_type(left.rsub(c), pd.Series), pd.Series)
67+
check(assert_type(left_i.rsub(b), pd.Series), pd.Series)
68+
check(assert_type(left_i.rsub(i), pd.Series), pd.Series)
69+
check(assert_type(left_i.rsub(f), pd.Series), pd.Series)
70+
check(assert_type(left_i.rsub(c), pd.Series), pd.Series)
6171

6272

6373
def test_sub_numpy_array() -> None:
@@ -67,35 +77,35 @@ def test_sub_numpy_array() -> None:
6777
f = np.array([1.0, 2.0, 3.0], np.float64)
6878
c = np.array([1.1j, 2.2j, 4.1j], np.complex128)
6979

70-
check(assert_type(left - b, pd.Series), pd.Series)
71-
check(assert_type(left - i, pd.Series), pd.Series)
72-
check(assert_type(left - f, pd.Series), pd.Series)
73-
check(assert_type(left - c, pd.Series), pd.Series)
80+
check(assert_type(left_i - b, pd.Series), pd.Series)
81+
check(assert_type(left_i - i, pd.Series), pd.Series)
82+
check(assert_type(left_i - f, pd.Series), pd.Series)
83+
check(assert_type(left_i - c, pd.Series), pd.Series)
7484

7585
# `numpy` typing gives the corresponding `ndarray`s in the static type
7686
# checking, where our `__rsub__` cannot override. At runtime, they return
7787
# `Series`s.
7888
# `mypy` thinks the return types are `Any`, which is a bug.
79-
check(assert_type(b - left, NoReturn), pd.Series) # type: ignore[assert-type]
89+
check(assert_type(b - left_i, NoReturn), pd.Series) # type: ignore[assert-type]
8090
check(
81-
assert_type(i - left, "npt.NDArray[np.int64]"), pd.Series # type: ignore[assert-type]
91+
assert_type(i - left_i, "npt.NDArray[np.int64]"), pd.Series # type: ignore[assert-type]
8292
)
8393
check(
84-
assert_type(f - left, "npt.NDArray[np.float64]"), pd.Series # type: ignore[assert-type]
94+
assert_type(f - left_i, "npt.NDArray[np.float64]"), pd.Series # type: ignore[assert-type]
8595
)
8696
check(
87-
assert_type(c - left, "npt.NDArray[np.complex128]"), pd.Series # type: ignore[assert-type]
97+
assert_type(c - left_i, "npt.NDArray[np.complex128]"), pd.Series # type: ignore[assert-type]
8898
)
8999

90-
check(assert_type(left.sub(b), pd.Series), pd.Series)
91-
check(assert_type(left.sub(i), pd.Series), pd.Series)
92-
check(assert_type(left.sub(f), pd.Series), pd.Series)
93-
check(assert_type(left.sub(c), pd.Series), pd.Series)
100+
check(assert_type(left_i.sub(b), pd.Series), pd.Series)
101+
check(assert_type(left_i.sub(i), pd.Series), pd.Series)
102+
check(assert_type(left_i.sub(f), pd.Series), pd.Series)
103+
check(assert_type(left_i.sub(c), pd.Series), pd.Series)
94104

95-
check(assert_type(left.rsub(b), pd.Series), pd.Series)
96-
check(assert_type(left.rsub(i), pd.Series), pd.Series)
97-
check(assert_type(left.rsub(f), pd.Series), pd.Series)
98-
check(assert_type(left.rsub(c), pd.Series), pd.Series)
105+
check(assert_type(left_i.rsub(b), pd.Series), pd.Series)
106+
check(assert_type(left_i.rsub(i), pd.Series), pd.Series)
107+
check(assert_type(left_i.rsub(f), pd.Series), pd.Series)
108+
check(assert_type(left_i.rsub(c), pd.Series), pd.Series)
99109

100110

101111
def test_sub_pd_series() -> None:
@@ -105,22 +115,78 @@ def test_sub_pd_series() -> None:
105115
f = pd.Series([1.0, 2.0, 3.0])
106116
c = pd.Series([1.1j, 2.2j, 4.1j])
107117

108-
check(assert_type(left - b, pd.Series), pd.Series)
109-
check(assert_type(left - i, pd.Series), pd.Series)
110-
check(assert_type(left - f, pd.Series), pd.Series)
111-
check(assert_type(left - c, pd.Series), pd.Series)
112-
113-
check(assert_type(b - left, pd.Series), pd.Series)
114-
check(assert_type(i - left, pd.Series), pd.Series)
115-
check(assert_type(f - left, pd.Series), pd.Series)
116-
check(assert_type(c - left, pd.Series), pd.Series)
117-
118-
check(assert_type(left.sub(b), pd.Series), pd.Series)
119-
check(assert_type(left.sub(i), pd.Series), pd.Series)
120-
check(assert_type(left.sub(f), pd.Series), pd.Series)
121-
check(assert_type(left.sub(c), pd.Series), pd.Series)
122-
123-
check(assert_type(left.rsub(b), pd.Series), pd.Series)
124-
check(assert_type(left.rsub(i), pd.Series), pd.Series)
125-
check(assert_type(left.rsub(f), pd.Series), pd.Series)
126-
check(assert_type(left.rsub(c), pd.Series), pd.Series)
118+
check(assert_type(left_i - b, pd.Series), pd.Series)
119+
check(assert_type(left_i - i, pd.Series), pd.Series)
120+
check(assert_type(left_i - f, pd.Series), pd.Series)
121+
check(assert_type(left_i - c, pd.Series), pd.Series)
122+
123+
check(assert_type(b - left_i, pd.Series), pd.Series)
124+
check(assert_type(i - left_i, pd.Series), pd.Series)
125+
check(assert_type(f - left_i, pd.Series), pd.Series)
126+
check(assert_type(c - left_i, pd.Series), pd.Series)
127+
128+
check(assert_type(left_i.sub(b), pd.Series), pd.Series)
129+
check(assert_type(left_i.sub(i), pd.Series), pd.Series)
130+
check(assert_type(left_i.sub(f), pd.Series), pd.Series)
131+
check(assert_type(left_i.sub(c), pd.Series), pd.Series)
132+
133+
check(assert_type(left_i.rsub(b), pd.Series), pd.Series)
134+
check(assert_type(left_i.rsub(i), pd.Series), pd.Series)
135+
check(assert_type(left_i.rsub(f), pd.Series), pd.Series)
136+
check(assert_type(left_i.rsub(c), pd.Series), pd.Series)
137+
138+
139+
anchor = datetime(2025, 8, 18)
140+
left_ts = pd.DataFrame({"a": [anchor + timedelta(hours=h + 1) for h in range(3)]})["a"]
141+
142+
143+
def test_sub_py_datetime() -> None:
144+
"""Test pd.Series[Any] - Python native datetime(s)"""
145+
s = anchor
146+
147+
check(assert_type(left_ts - s, "TimedeltaSeries"), pd.Series, pd.Timedelta)
148+
149+
check(assert_type(s - left_ts, "TimedeltaSeries"), pd.Series, pd.Timedelta)
150+
151+
check(assert_type(left_ts.sub(s), "TimedeltaSeries"), pd.Series, pd.Timedelta)
152+
153+
check(assert_type(left_ts.rsub(s), "TimedeltaSeries"), pd.Series, pd.Timedelta)
154+
155+
156+
def test_sub_numpy_datetime() -> None:
157+
"""Test pd.Series[Any] - numpy datetime(s)"""
158+
s = np.datetime64(anchor)
159+
a = np.array([s + np.timedelta64(m, "m") for m in range(3)])
160+
161+
# `numpy` typing gives the corresponding `ndarray`s in the static type
162+
# checking, where our `__rsub__` cannot override. At runtime, they return
163+
# `Series`s.
164+
check(assert_type(left_ts - s, "TimedeltaSeries"), pd.Series, pd.Timedelta)
165+
check(assert_type(left_ts - a, "TimedeltaSeries"), pd.Series, pd.Timedelta) # type: ignore[assert-type]
166+
167+
check(assert_type(s - left_ts, "TimedeltaSeries"), pd.Series, pd.Timedelta)
168+
check(assert_type(a - left_ts, "npt.NDArray[np.datetime64]"), pd.Series, pd.Timedelta) # type: ignore[assert-type]
169+
170+
check(assert_type(left_ts.sub(s), "TimedeltaSeries"), pd.Series, pd.Timedelta)
171+
check(assert_type(left_ts.sub(a), "TimedeltaSeries"), pd.Series, pd.Timedelta) # type: ignore[assert-type]
172+
173+
check(assert_type(left_ts.rsub(s), "TimedeltaSeries"), pd.Series, pd.Timedelta)
174+
check(assert_type(left_ts.rsub(a), "TimedeltaSeries"), pd.Series, pd.Timedelta) # type: ignore[assert-type]
175+
176+
177+
def test_sub_pd_datetime() -> None:
178+
"""Test pd.Series[Any] - Pandas datetime(s)"""
179+
s = pd.Timestamp(anchor)
180+
a = pd.Series([s + pd.Timedelta(minutes=m) for m in range(3)])
181+
182+
check(assert_type(left_ts - s, "TimedeltaSeries"), pd.Series, pd.Timedelta)
183+
check(assert_type(left_ts - a, "TimedeltaSeries"), pd.Series, pd.Timedelta)
184+
185+
check(assert_type(s - left_ts, "TimedeltaSeries"), pd.Series, pd.Timedelta)
186+
check(assert_type(a - left_ts, "TimedeltaSeries"), pd.Series, pd.Timedelta)
187+
188+
check(assert_type(left_ts.sub(s), "TimedeltaSeries"), pd.Series, pd.Timedelta)
189+
check(assert_type(left_ts.sub(a), "TimedeltaSeries"), pd.Series, pd.Timedelta)
190+
191+
check(assert_type(left_ts.rsub(s), "TimedeltaSeries"), pd.Series, pd.Timedelta)
192+
check(assert_type(left_ts.rsub(a), "TimedeltaSeries"), pd.Series, pd.Timedelta)

0 commit comments

Comments
 (0)