Skip to content

Commit 6469196

Browse files
committed
feat: tests for addition
1 parent 7a87d18 commit 6469196

File tree

8 files changed

+238
-16
lines changed

8 files changed

+238
-16
lines changed

attempt.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
1-
from typing import assert_type, reveal_type
1+
import datetime as dt
2+
from typing import (
3+
assert_type,
4+
reveal_type,
5+
)
6+
7+
import numpy as np
28
import pandas as pd
39
from pandas.core.series import TimedeltaSeries # noqa: F401
4-
import numpy as np
5-
import datetime as dt
610

711
from tests import check
812

9-
1013
df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [5, 4, 3, 2, 1]})
1114
s1 = df.min(axis=1)
1215
s2 = df.max(axis=1)
@@ -45,4 +48,4 @@
4548
r7 = ts1 - tsp1
4649
check(assert_type(r7, "TimedeltaSeries"), pd.Series, pd.Timedelta)
4750
r8 = ts1 - dt1
48-
check(assert_type(r8, "TimedeltaSeries"), pd.Series, pd.Timedelta)
51+
check(assert_type(r8, "TimedeltaSeries"), pd.Series, pd.Timedelta)

pandas-stubs/core/series.pyi

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,7 @@ from pandas._typing import (
187187
np_ndarray_complex,
188188
np_ndarray_dt,
189189
np_ndarray_float,
190+
np_ndarray_str,
190191
np_ndarray_td,
191192
npt,
192193
num,
@@ -1688,9 +1689,11 @@ class Series(IndexOpsMixin[S1], NDFrame):
16881689
) -> Series[complex]: ...
16891690
@overload
16901691
def __add__(
1691-
self: Series[_str], other: _str | Sequence[_str] | Series[_str]
1692+
self: Series[_str], other: _str | Sequence[_str] | np_ndarray_str | Series[_str]
16921693
) -> Series[_str]: ...
16931694
@overload
1695+
def __add__(self: Series[Timestamp], other: np_ndarray_dt) -> Never: ...
1696+
@overload
16941697
def __add__(
16951698
self: Series[Timestamp],
16961699
other: (
@@ -1842,7 +1845,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
18421845
@overload
18431846
def add(
18441847
self: Series[_str],
1845-
other: _str | Sequence[_str] | Series[_str],
1848+
other: _str | Sequence[_str] | np_ndarray_str | Series[_str],
18461849
level: Level | None = None,
18471850
fill_value: float | None = None,
18481851
axis: int = 0,
@@ -1949,9 +1952,11 @@ class Series(IndexOpsMixin[S1], NDFrame):
19491952
) -> Series[complex]: ...
19501953
@overload
19511954
def __radd__(
1952-
self: Series[_str], other: _str | Sequence[_str] | Series[_str]
1955+
self: Series[_str], other: _str | Sequence[_str] | np_ndarray_str | Series[_str]
19531956
) -> Series[_str]: ...
19541957
@overload
1958+
def __radd__(self: Series[Timestamp], other: np_ndarray_dt) -> Never: ...
1959+
@overload
19551960
def __radd__(
19561961
self: Series[Timestamp],
19571962
other: (
@@ -2094,7 +2099,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
20942099
@overload
20952100
def radd(
20962101
self: Series[_str],
2097-
other: _str | Sequence[_str] | Series[_str],
2102+
other: _str | Sequence[_str] | np_ndarray_str | Series[_str],
20982103
level: Level | None = None,
20992104
fill_value: float | None = None,
21002105
axis: int = 0,

tests/series/arithmetic/str/__init__.py

Whitespace-only changes.
Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
from typing import Any
2+
3+
import numpy as np
4+
import pandas as pd
5+
from typing_extensions import assert_type
6+
7+
from tests import check
8+
9+
left = pd.Series(["1", "23", "456"]) # left operand
10+
11+
12+
def test_add_py_scalar() -> None:
13+
"""Test pd.Series[str] + Python native str"""
14+
r0 = "right"
15+
16+
check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
17+
18+
check(assert_type(r0 + left, "pd.Series[str]"), pd.Series, str)
19+
20+
check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str)
21+
22+
check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str)
23+
24+
25+
def test_add_py_sequence() -> None:
26+
"""Test pd.Series[str] + Python native sequence"""
27+
r0 = ["a", "bc", "def"]
28+
r1 = tuple(r0)
29+
30+
check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
31+
check(assert_type(left + r1, "pd.Series[str]"), pd.Series, str)
32+
33+
check(assert_type(r0 + left, "pd.Series[str]"), pd.Series, str)
34+
check(assert_type(r1 + left, "pd.Series[str]"), pd.Series, str)
35+
36+
check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str)
37+
check(assert_type(left.add(r1), "pd.Series[str]"), pd.Series, str)
38+
39+
check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str)
40+
check(assert_type(left.radd(r1), "pd.Series[str]"), pd.Series, str)
41+
42+
43+
def test_add_numpy_array() -> None:
44+
"""Test pd.Series[str] + numpy array"""
45+
r0 = np.array(["a", "bc", "def"], np.str_)
46+
47+
check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
48+
49+
# `numpy` typing gives `Any` in the static type
50+
# checking, where our `__radd__` cannot override. At runtime, they return
51+
# `Series`s.
52+
check(assert_type(r0 + left, Any), pd.Series, str)
53+
54+
check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str)
55+
56+
check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str)
57+
58+
59+
def test_add_pd_series() -> None:
60+
"""Test pd.Series[str] + pandas series"""
61+
r0 = pd.Series(["a", "bc", "def"])
62+
63+
check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
64+
65+
check(assert_type(r0 + left, "pd.Series[str]"), pd.Series, str)
66+
67+
check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str)
68+
69+
check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str)

tests/series/arithmetic/timestamp/__init__.py

Whitespace-only changes.
Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
from datetime import (
2+
datetime,
3+
timedelta,
4+
)
5+
6+
import numpy as np
7+
from numpy import typing as npt # noqa: F401
8+
import pandas as pd
9+
from typing_extensions import (
10+
Never,
11+
assert_type,
12+
)
13+
14+
from tests import (
15+
TYPE_CHECKING_INVALID_USAGE,
16+
check,
17+
)
18+
19+
left = pd.Series([pd.Timestamp(2025, 8, 20)])
20+
21+
22+
def test_add_py_scalar() -> None:
23+
"""Test pd.Series[pd.Timestamp] + Python native scalars"""
24+
s = datetime(2025, 8, 20)
25+
d = timedelta(seconds=1)
26+
27+
if TYPE_CHECKING_INVALID_USAGE:
28+
_0 = left + s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
29+
check(assert_type(left + d, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
30+
31+
if TYPE_CHECKING_INVALID_USAGE:
32+
_1 = s + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
33+
check(assert_type(d + left, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
34+
35+
if TYPE_CHECKING_INVALID_USAGE:
36+
left.add(s) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
37+
check(assert_type(left.add(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
38+
39+
if TYPE_CHECKING_INVALID_USAGE:
40+
left.radd(s) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
41+
check(assert_type(left.radd(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
42+
43+
44+
def test_add_numpy_scalar() -> None:
45+
"""Test pd.Series[pd.Timestamp] + numpy scalars"""
46+
s = np.datetime64("2025-08-20")
47+
d = np.timedelta64(1, "s")
48+
49+
if TYPE_CHECKING_INVALID_USAGE:
50+
_0 = left + s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
51+
check(assert_type(left + d, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
52+
53+
if TYPE_CHECKING_INVALID_USAGE:
54+
_1 = s + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
55+
check(assert_type(d + left, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
56+
57+
if TYPE_CHECKING_INVALID_USAGE:
58+
left.add(s) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
59+
check(assert_type(left.add(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
60+
61+
if TYPE_CHECKING_INVALID_USAGE:
62+
left.radd(s) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
63+
check(assert_type(left.radd(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
64+
65+
66+
def test_add_pd_scalar() -> None:
67+
"""Test pd.Series[pd.Timestamp] + pandas scalars"""
68+
s = pd.Timestamp("2025-08-20")
69+
d = pd.Timedelta(seconds=1)
70+
71+
if TYPE_CHECKING_INVALID_USAGE:
72+
_0 = left + s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
73+
check(assert_type(left + d, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
74+
75+
if TYPE_CHECKING_INVALID_USAGE:
76+
_1 = s + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
77+
check(assert_type(d + left, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
78+
79+
if TYPE_CHECKING_INVALID_USAGE:
80+
left.add(s) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
81+
check(assert_type(left.add(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
82+
83+
if TYPE_CHECKING_INVALID_USAGE:
84+
left.radd(s) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
85+
check(assert_type(left.radd(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
86+
87+
88+
def test_add_py_sequence() -> None:
89+
"""Test pd.Series[pd.Timestamp] + Python native sequence"""
90+
s = [datetime(2025, 8, 20)]
91+
d = [timedelta(seconds=1)]
92+
93+
if TYPE_CHECKING_INVALID_USAGE:
94+
_0 = left + s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
95+
_a = left + d # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
96+
97+
_1 = s + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
98+
_b = d + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
99+
100+
left.add(s) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
101+
left.add(d) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
102+
103+
left.radd(s) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
104+
left.radd(d) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
105+
106+
107+
def test_add_numpy_array() -> None:
108+
"""Test pd.Series[pd.Timestamp] + numpy array"""
109+
s = np.array([np.datetime64("2025-08-20")], np.datetime64)
110+
d = np.array([np.timedelta64(1, "s")], np.timedelta64)
111+
112+
if TYPE_CHECKING_INVALID_USAGE:
113+
assert_type(left + s, Never)
114+
check(assert_type(left + d, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
115+
116+
# `numpy` typing gives the corresponding `ndarray`s in the static type
117+
# checking, where our `__radd__` cannot override. At runtime, they return
118+
# `Series`s.
119+
if TYPE_CHECKING_INVALID_USAGE:
120+
assert_type(s + left, "npt.NDArray[np.datetime64]")
121+
check(assert_type(d + left, "npt.NDArray[np.timedelta64]"), pd.Series, pd.Timestamp)
122+
123+
if TYPE_CHECKING_INVALID_USAGE:
124+
left.add(s) # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
125+
check(assert_type(left.add(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
126+
127+
if TYPE_CHECKING_INVALID_USAGE:
128+
left.radd(s) # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
129+
check(assert_type(left.radd(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
130+
131+
132+
def test_add_pd_series() -> None:
133+
"""Test pd.Series[pd.Timestamp] + pandas Series"""
134+
s = pd.Series([pd.Timestamp("2025-08-20")])
135+
d = pd.Series([pd.Timedelta(seconds=1)])
136+
137+
if TYPE_CHECKING_INVALID_USAGE:
138+
_0 = left + s # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
139+
check(assert_type(left + d, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
140+
141+
if TYPE_CHECKING_INVALID_USAGE:
142+
_1 = s + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
143+
check(assert_type(d + left, "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
144+
145+
if TYPE_CHECKING_INVALID_USAGE:
146+
left.add(s) # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
147+
check(assert_type(left.add(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)
148+
149+
if TYPE_CHECKING_INVALID_USAGE:
150+
left.radd(s) # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
151+
check(assert_type(left.radd(d), "pd.Series[pd.Timestamp]"), pd.Series, pd.Timestamp)

tests/series/test_series.py

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1653,12 +1653,6 @@ def test_reset_index() -> None:
16531653
assert assert_type(s.reset_index(inplace=True, drop=True), None) is None
16541654

16551655

1656-
def test_series_add_str() -> None:
1657-
s = pd.Series(["abc", "def"])
1658-
check(assert_type(s + "x", "pd.Series[str]"), pd.Series, str)
1659-
check(assert_type("x" + s, "pd.Series[str]"), pd.Series, str)
1660-
1661-
16621656
def test_series_dtype() -> None:
16631657
s = pd.Series(["abc", "def"], dtype=str)
16641658
check(assert_type(s, "pd.Series[str]"), pd.Series, str)

tests/test_timefuncs.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1836,7 +1836,7 @@ def test_timestamp_to_list_add() -> None:
18361836
sseries + pd.Timedelta(1, "d")
18371837

18381838
check(
1839-
assert_type(sseries + pd.Timedelta(1, "D"), pd.Series[pd.Timestamp]),
1839+
assert_type(sseries + pd.Timedelta(1, "D"), "pd.Series[pd.Timestamp]"),
18401840
pd.Series,
18411841
pd.Timestamp,
18421842
)

0 commit comments

Comments
 (0)