Skip to content

Commit 480aefd

Browse files
committed
feat(series): addition for str
1 parent 766cc1d commit 480aefd

File tree

4 files changed

+162
-19
lines changed

4 files changed

+162
-19
lines changed

pandas-stubs/core/series.pyi

Lines changed: 39 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,7 @@ from pandas._typing import (
188188
np_ndarray_complex,
189189
np_ndarray_dt,
190190
np_ndarray_float,
191+
np_ndarray_str,
191192
np_ndarray_td,
192193
npt,
193194
num,
@@ -1673,7 +1674,8 @@ class Series(IndexOpsMixin[S1], NDFrame):
16731674
def __add__(
16741675
self: Series[complex],
16751676
other: (
1676-
Sequence[_T_COMPLEX]
1677+
_T_COMPLEX
1678+
| Sequence[_T_COMPLEX]
16771679
| np_ndarray_bool
16781680
| np_ndarray_anyint
16791681
| np_ndarray_float
@@ -1682,7 +1684,16 @@ class Series(IndexOpsMixin[S1], NDFrame):
16821684
),
16831685
) -> Series[complex]: ...
16841686
@overload
1685-
def __add__(self, other: S1 | Series[S1]) -> Self: ...
1687+
def __add__(
1688+
self: Series[_str],
1689+
other: (
1690+
np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex
1691+
),
1692+
) -> Never: ...
1693+
@overload
1694+
def __add__(
1695+
self: Series[_str], other: _str | Sequence[_str] | np_ndarray_str | Series[_str]
1696+
) -> Series[_str]: ...
16861697
@overload
16871698
def add(
16881699
self: Series[Never],
@@ -1808,7 +1819,8 @@ class Series(IndexOpsMixin[S1], NDFrame):
18081819
def add(
18091820
self: Series[complex],
18101821
other: (
1811-
Sequence[_T_COMPLEX]
1822+
_T_COMPLEX
1823+
| Sequence[_T_COMPLEX]
18121824
| np_ndarray_bool
18131825
| np_ndarray_anyint
18141826
| np_ndarray_float
@@ -1821,13 +1833,13 @@ class Series(IndexOpsMixin[S1], NDFrame):
18211833
) -> Series[complex]: ...
18221834
@overload
18231835
def add(
1824-
self,
1825-
other: S1 | Series[S1],
1836+
self: Series[_str],
1837+
other: _str | Sequence[_str] | np_ndarray_str | Series[_str],
18261838
level: Level | None = None,
18271839
fill_value: float | None = None,
18281840
axis: int = 0,
1829-
) -> Self: ...
1830-
@overload
1841+
) -> Series[_str]: ...
1842+
@overload # type: ignore[override]
18311843
def __radd__(self: Series[Never], other: Scalar | _ListLike) -> Series: ...
18321844
@overload
18331845
def __radd__(
@@ -1874,7 +1886,8 @@ class Series(IndexOpsMixin[S1], NDFrame):
18741886
def __radd__(
18751887
self: Series[complex],
18761888
other: (
1877-
Sequence[_T_COMPLEX]
1889+
_T_COMPLEX
1890+
| Sequence[_T_COMPLEX]
18781891
| np_ndarray_bool
18791892
| np_ndarray_anyint
18801893
| np_ndarray_float
@@ -1886,7 +1899,16 @@ class Series(IndexOpsMixin[S1], NDFrame):
18861899
self: Series[_T_COMPLEX], other: np_ndarray_complex
18871900
) -> Series[complex]: ...
18881901
@overload
1889-
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
1902+
def __radd__(
1903+
self: Series[_str],
1904+
other: (
1905+
np_ndarray_bool | np_ndarray_anyint | np_ndarray_float | np_ndarray_complex
1906+
),
1907+
) -> Never: ...
1908+
@overload
1909+
def __radd__(
1910+
self: Series[_str], other: _str | Sequence[_str] | np_ndarray_str | Series[_str]
1911+
) -> Series[_str]: ...
18901912
@overload
18911913
def radd(
18921914
self: Series[Never],
@@ -1980,7 +2002,8 @@ class Series(IndexOpsMixin[S1], NDFrame):
19802002
def radd(
19812003
self: Series[complex],
19822004
other: (
1983-
Sequence[_T_COMPLEX]
2005+
_T_COMPLEX
2006+
| Sequence[_T_COMPLEX]
19842007
| np_ndarray_bool
19852008
| np_ndarray_anyint
19862009
| np_ndarray_float
@@ -2000,12 +2023,12 @@ class Series(IndexOpsMixin[S1], NDFrame):
20002023
) -> Series[complex]: ...
20012024
@overload
20022025
def radd(
2003-
self,
2004-
other: S1 | Series[S1],
2026+
self: Series[_str],
2027+
other: _str | Sequence[_str] | np_ndarray_str | Series[_str],
20052028
level: Level | None = None,
20062029
fill_value: float | None = None,
20072030
axis: int = 0,
2008-
) -> Self: ...
2031+
) -> Series[_str]: ...
20092032
# ignore needed for mypy as we want different results based on the arguments
20102033
@overload # type: ignore[override]
20112034
def __and__( # pyright: ignore[reportOverlappingOverload]
@@ -3682,7 +3705,7 @@ class TimestampSeries(_SeriesSubclassBase[Timestamp, np.datetime64]):
36823705
@property
36833706
def dt(self) -> TimestampProperties: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
36843707
def __add__(self, other: TimedeltaSeries | np.timedelta64 | timedelta | BaseOffset) -> TimestampSeries: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
3685-
def __radd__(self, other: TimedeltaSeries | np.timedelta64 | timedelta) -> TimestampSeries: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
3708+
def __radd__(self, other: TimedeltaSeries | np.timedelta64 | timedelta) -> TimestampSeries: ... # type: ignore[override]
36863709
@overload # type: ignore[override]
36873710
def __sub__(
36883711
self, other: Timestamp | datetime | TimestampSeries
@@ -3742,7 +3765,7 @@ class TimedeltaSeries(_SeriesSubclassBase[Timedelta, np.timedelta64]):
37423765
def __add__( # pyright: ignore[reportIncompatibleMethodOverride]
37433766
self, other: timedelta | Timedelta | np.timedelta64
37443767
) -> TimedeltaSeries: ...
3745-
def __radd__(self, other: datetime | Timestamp | TimestampSeries) -> TimestampSeries: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
3768+
def __radd__(self, other: datetime | Timestamp | TimestampSeries) -> TimestampSeries: ... # type: ignore[override]
37463769
def __mul__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
37473770
self, other: num | Sequence[num] | Series[int] | Series[float]
37483771
) -> TimedeltaSeries: ...
@@ -3859,9 +3882,7 @@ class OffsetSeries(_SeriesSubclassBase[BaseOffset, np.object_]):
38593882
@overload # type: ignore[override]
38603883
def __radd__(self, other: Period) -> PeriodSeries: ...
38613884
@overload
3862-
def __radd__( # pyright: ignore[reportIncompatibleMethodOverride]
3863-
self, other: BaseOffset
3864-
) -> OffsetSeries: ...
3885+
def __radd__(self, other: BaseOffset) -> OffsetSeries: ...
38653886
def cumprod(
38663887
self,
38673888
axis: AxisIndex = ...,

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ mypy = "1.17.0"
3939
pandas = "2.3.1"
4040
pyarrow = ">=10.0.1"
4141
pytest = ">=7.1.2"
42-
pyright = ">=1.1.403"
42+
pyright = ">=1.1.404"
4343
ty = "^0.0.1a8"
4444
pyrefly = "^0.21.0"
4545
poethepoet = ">=0.16.5"

tests/series/arithmetic/str/__init__.py

Whitespace-only changes.
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
import sys
2+
from typing import Any
3+
4+
import numpy as np
5+
from numpy import typing as npt # noqa: F401
6+
import pandas as pd
7+
from typing_extensions import (
8+
Never,
9+
assert_type,
10+
)
11+
12+
from tests import (
13+
TYPE_CHECKING_INVALID_USAGE,
14+
check,
15+
)
16+
17+
left = pd.Series(["1", "23", "456"]) # left operand
18+
19+
20+
def test_add_py_scalar() -> None:
21+
"""Testpd.Series[str]+ Python native 'scalar's"""
22+
i = 4
23+
r0 = "right"
24+
25+
if TYPE_CHECKING_INVALID_USAGE:
26+
_0 = left + i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
27+
check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
28+
29+
if TYPE_CHECKING_INVALID_USAGE:
30+
_1 = i + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
31+
check(assert_type(r0 + left, "pd.Series[str]"), pd.Series, str)
32+
33+
if TYPE_CHECKING_INVALID_USAGE:
34+
left.add(i) # type: ignore[call-overload] # pyright: ignore[reportArgumentType,reportCallIssue]
35+
check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str)
36+
37+
if TYPE_CHECKING_INVALID_USAGE:
38+
left.radd(i) # type: ignore[call-overload] # pyright: ignore[reportArgumentType]
39+
check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str)
40+
41+
42+
def test_add_py_sequence() -> None:
43+
"""Testpd.Series[str]+ Python native sequence"""
44+
i = [3, 5, 8]
45+
r0 = ["a", "bc", "def"]
46+
r1 = tuple(r0)
47+
48+
if TYPE_CHECKING_INVALID_USAGE:
49+
_0 = left + i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
50+
check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
51+
check(assert_type(left + r1, "pd.Series[str]"), pd.Series, str)
52+
53+
if TYPE_CHECKING_INVALID_USAGE:
54+
_1 = i + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
55+
check(assert_type(r0 + left, "pd.Series[str]"), pd.Series, str)
56+
check(assert_type(r1 + left, "pd.Series[str]"), pd.Series, str)
57+
58+
if TYPE_CHECKING_INVALID_USAGE:
59+
left.add(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
60+
check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str)
61+
check(assert_type(left.add(r1), "pd.Series[str]"), pd.Series, str)
62+
63+
if TYPE_CHECKING_INVALID_USAGE:
64+
left.radd(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
65+
check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str)
66+
check(assert_type(left.radd(r1), "pd.Series[str]"), pd.Series, str)
67+
68+
69+
def test_add_numpy_array() -> None:
70+
"""Testpd.Series[str]+ numpy array"""
71+
i = np.array([3, 5, 8], np.int64)
72+
r0 = np.array(["a", "bc", "def"], np.str_)
73+
74+
if TYPE_CHECKING_INVALID_USAGE:
75+
assert_type(left + i, Never)
76+
check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
77+
78+
# `numpy` typing gives `npt.NDArray[np.int64]` in the static type
79+
# checking, where our `__radd__` cannot override. At runtime, they return
80+
# `Series`s.
81+
if TYPE_CHECKING_INVALID_USAGE:
82+
assert_type(i + left, "npt.NDArray[np.int64]")
83+
if sys.version_info >= (3, 11):
84+
# `numpy` typing gives `npt.NDArray[np.int64]` in the static type
85+
# checking, where our `__radd__` cannot override. At runtime, they return
86+
# `Series`s.
87+
check(assert_type(r0 + left, "npt.NDArray[np.str_]"), pd.Series, str)
88+
else:
89+
# Python 3.10 uses NumPy 2.2.6, and it has for r0 ndarray[tuple[int,...], dtype[str_]]
90+
# Python 3.11+ uses NumPy 2.3.2, and it has for r0 ndarray[tuple[Any,...,dtype[str_]]
91+
# https://github.com/pandas-dev/pandas-stubs/pull/1274#discussion_r2291498975
92+
check(assert_type(r0 + left, Any), pd.Series, str)
93+
94+
if TYPE_CHECKING_INVALID_USAGE:
95+
left.add(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
96+
check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str)
97+
98+
if TYPE_CHECKING_INVALID_USAGE:
99+
left.radd(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
100+
check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str)
101+
102+
103+
def test_add_pd_series() -> None:
104+
"""Testpd.Series[str]+ pandas series"""
105+
i = pd.Series([3, 5, 8])
106+
r0 = pd.Series(["a", "bc", "def"])
107+
108+
if TYPE_CHECKING_INVALID_USAGE:
109+
_0 = left + i # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
110+
check(assert_type(left + r0, "pd.Series[str]"), pd.Series, str)
111+
112+
if TYPE_CHECKING_INVALID_USAGE:
113+
_1 = i + left # type: ignore[operator] # pyright: ignore[reportOperatorIssue]
114+
check(assert_type(r0 + left, "pd.Series[str]"), pd.Series, str)
115+
116+
if TYPE_CHECKING_INVALID_USAGE:
117+
left.add(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType,reportCallIssue]
118+
check(assert_type(left.add(r0), "pd.Series[str]"), pd.Series, str)
119+
120+
if TYPE_CHECKING_INVALID_USAGE:
121+
left.radd(i) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
122+
check(assert_type(left.radd(r0), "pd.Series[str]"), pd.Series, str)

0 commit comments

Comments
 (0)