Skip to content

Commit c494518

Browse files
cmp0xffcmp0xff
committed
feat: Series[int] +
1 parent c13e5cd commit c494518

File tree

5 files changed

+191
-108
lines changed

5 files changed

+191
-108
lines changed

‎pandas-stubs/core/series.pyi

Lines changed: 127 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ from typing import (
2525
Generic,
2626
Literal,
2727
NoReturn,
28+
TypeVar,
2829
final,
2930
overload,
3031
)
@@ -186,6 +187,8 @@ from pandas.core.dtypes.dtypes import CategoricalDtype
186187

187188
from pandas.plotting import PlotAccessor
188189

190+
_T_COMPLEX = TypeVar("_T_COMPLEX", bound=complex)
191+
189192
class _iLocIndexerSeries(_iLocIndexer, Generic[S1]):
190193
# get item
191194
@overload
@@ -1620,8 +1623,34 @@ class Series(IndexOpsMixin[S1], NDFrame):
16201623
# them up.
16211624
@overload
16221625
def __add__(
1623-
self: Series[int], other: int | Sequence[int] | Series[int]
1624-
) -> Series[int]: ...
1626+
self: Series[Never],
1627+
other: num | _str | timedelta | Timedelta | _ListLike | Series | np.timedelta64,
1628+
) -> Series: ...
1629+
@overload
1630+
def __add__(self: Series[int], other: Series[Never]) -> Series: ...
1631+
@overload
1632+
def __add__(
1633+
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX]
1634+
) -> Series[_T_COMPLEX]: ...
1635+
@overload
1636+
def __add__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ...
1637+
@overload
1638+
def __add__(self: Series[int], other: np_ndarray_float) -> Series[float]: ...
1639+
@overload
1640+
def __add__(self: Series[int], other: np_ndarray_complex) -> Series[complex]: ...
1641+
@overload
1642+
def __add__(self: Series[float], other: Series[Never]) -> Series: ...
1643+
@overload
1644+
def __add__(
1645+
self: Series[float],
1646+
other: int | Sequence[int] | np_ndarray_anyint | Series[int],
1647+
) -> Series[float]: ...
1648+
@overload
1649+
def __add__(self: Series[complex], other: Series[Never]) -> Series: ...
1650+
@overload
1651+
def __add__(
1652+
self: Series[complex], other: int | Sequence[int] | Series[int]
1653+
) -> Series[complex]: ...
16251654
@overload
16261655
def __add__(self, other: S1 | Self) -> Self: ...
16271656
@overload
@@ -1669,6 +1698,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
16691698
@overload
16701699
def __or__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
16711700
@overload
1701+
def __radd__(
1702+
self: Series[Never], other: num | _str | _ListLike | Series
1703+
) -> Series: ...
1704+
@overload
1705+
def __radd__(
1706+
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1707+
) -> Series[_T_COMPLEX]: ...
1708+
@overload
16721709
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
16731710
@overload
16741711
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1747,19 +1784,60 @@ class Series(IndexOpsMixin[S1], NDFrame):
17471784
@property
17481785
def loc(self) -> _LocIndexerSeries[S1]: ...
17491786
# Methods
1787+
@overload
17501788
def add(
1751-
self,
1789+
self: Series[Never],
17521790
other: (
1753-
Sequence[S1]
1791+
Sequence
17541792
| np_ndarray_anyint
17551793
| np_ndarray_float
17561794
| np_ndarray_complex
1757-
| Series[S1]
1795+
| Series
17581796
| Scalar
17591797
),
17601798
level: Level | None = ...,
17611799
fill_value: float | None = ...,
17621800
axis: int = ...,
1801+
) -> Series: ...
1802+
@overload
1803+
def add(
1804+
self: Series[int],
1805+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1806+
level: Level | None = ...,
1807+
fill_value: float | None = ...,
1808+
axis: int = ...,
1809+
) -> Series[_T_COMPLEX]: ...
1810+
@overload
1811+
def add(
1812+
self: Series[int],
1813+
other: np_ndarray_anyint,
1814+
level: Level | None = ...,
1815+
fill_value: float | None = ...,
1816+
axis: int = ...,
1817+
) -> Series[int]: ...
1818+
@overload
1819+
def add(
1820+
self: Series[int],
1821+
other: np_ndarray_float,
1822+
level: Level | None = ...,
1823+
fill_value: float | None = ...,
1824+
axis: int = ...,
1825+
) -> Series[float]: ...
1826+
@overload
1827+
def add(
1828+
self: Series[int],
1829+
other: np_ndarray_complex,
1830+
level: Level | None = ...,
1831+
fill_value: float | None = ...,
1832+
axis: int = ...,
1833+
) -> Series[complex]: ...
1834+
@overload
1835+
def add(
1836+
self,
1837+
other: Series[S1] | Scalar,
1838+
level: Level | None = ...,
1839+
fill_value: float | None = ...,
1840+
axis: int = ...,
17631841
) -> Series[S1]: ...
17641842
def all(
17651843
self,
@@ -1998,19 +2076,60 @@ class Series(IndexOpsMixin[S1], NDFrame):
19982076
min_count: int = ...,
19992077
**kwargs: Any,
20002078
) -> Scalar: ...
2079+
@overload
20012080
def radd(
2002-
self,
2081+
self: Series[Never],
20032082
other: (
2004-
Sequence[S1]
2083+
Sequence
20052084
| np_ndarray_anyint
20062085
| np_ndarray_float
20072086
| np_ndarray_complex
2008-
| Series[S1]
2087+
| Series
20092088
| Scalar
20102089
),
20112090
level: Level | None = ...,
20122091
fill_value: float | None = ...,
20132092
axis: AxisIndex = ...,
2093+
) -> Series: ...
2094+
@overload
2095+
def radd(
2096+
self: Series[int],
2097+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
2098+
level: Level | None = ...,
2099+
fill_value: float | None = ...,
2100+
axis: int = ...,
2101+
) -> Series[_T_COMPLEX]: ...
2102+
@overload
2103+
def radd(
2104+
self: Series[int],
2105+
other: np_ndarray_anyint,
2106+
level: Level | None = ...,
2107+
fill_value: float | None = ...,
2108+
axis: int = ...,
2109+
) -> Series[int]: ...
2110+
@overload
2111+
def radd(
2112+
self: Series[int],
2113+
other: np_ndarray_float,
2114+
level: Level | None = ...,
2115+
fill_value: float | None = ...,
2116+
axis: int = ...,
2117+
) -> Series[float]: ...
2118+
@overload
2119+
def radd(
2120+
self: Series[int],
2121+
other: np_ndarray_complex,
2122+
level: Level | None = ...,
2123+
fill_value: float | None = ...,
2124+
axis: int = ...,
2125+
) -> Series[complex]: ...
2126+
@overload
2127+
def radd(
2128+
self,
2129+
other: Series[S1] | Scalar,
2130+
level: Level | None = ...,
2131+
fill_value: float | None = ...,
2132+
axis: AxisIndex = ...,
20142133
) -> Series[S1]: ...
20152134
def rdivmod(
20162135
self,

‎tests/series/arithmetic/int/test_add.py

Lines changed: 56 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from typing import assert_type
22

3+
import numpy as np
34
import pandas as pd
45

56
from tests import check
@@ -11,94 +12,85 @@ def test_add_py_scalar() -> None:
1112
"""Test pd.Series[int] + Python native scalars"""
1213
i, f, c = 1, 1.0, 1j
1314

14-
check(assert_type(left + i, "pd.Series[int]"), pd.Series, int)
15-
# check(assert_type(left + f, "pd.Series[float]"), pd.Series, float)
16-
# check(assert_type(left + c, "pd.Series[complex]"), pd.Series, complex)
15+
check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.int64)
16+
check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64)
17+
check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128)
1718

18-
check(assert_type(i + left, "pd.Series[int]"), pd.Series, int)
19-
# check(assert_type(f + left, "pd.Series[float]"), pd.Series, float)
20-
# check(assert_type(c + left, "pd.Series[complex]"), pd.Series, complex)
19+
check(assert_type(i + left, "pd.Series[int]"), pd.Series, np.int64)
20+
check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64)
21+
check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128)
2122

22-
check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, int)
23-
# check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, float)
24-
# check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, complex)
23+
check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.int64)
24+
check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64)
25+
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)
2526

26-
check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, int)
27-
28-
29-
# check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, float)
30-
# check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, complex)
27+
check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.int64)
28+
check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64)
29+
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)
3130

3231

3332
def test_add_py_sequence() -> None:
3433
"""Test pd.Series[int] + Python native sequence"""
3534
i, f, c = [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
3635

37-
check(assert_type(left + i, "pd.Series[int]"), pd.Series, int)
38-
# check(assert_type(left + f, "pd.Series[float]"), pd.Series, float)
39-
# check(assert_type(left + c, "pd.Series[complex]"), pd.Series, complex)
40-
41-
# check(assert_type(i + left, "pd.Series[int]"), pd.Series, int)
42-
# check(assert_type(f + left, "pd.Series[float]"), pd.Series, float)
43-
# check(assert_type(c + left, "pd.Series[complex]"), pd.Series, complex)
36+
check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.int64)
37+
check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64)
38+
check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128)
4439

45-
check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, int)
40+
check(assert_type(i + left, "pd.Series[int]"), pd.Series, np.int64)
41+
check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64)
42+
check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128)
4643

44+
check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.int64)
45+
check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64)
46+
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)
4747

48-
# check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, float)
49-
# check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, complex)
48+
check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.int64)
49+
check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64)
50+
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)
5051

51-
# check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, int)
52-
# check(assert_type(left.radd(f), "pd.Series[int]"), pd.Series, int)
53-
# check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, complex)
5452

53+
def test_add_numpy_array() -> None:
54+
"""Test pd.Series[int] + numpy array"""
55+
i = np.array([2, 3, 5], np.int64)
56+
f = np.array([1.0, 2.0, 3.0], np.float64)
57+
c = np.array([1.1j, 2.2j, 4.1j], np.complex128)
5558

56-
# def test_add_numpy_array() -> None:
57-
# """Test pd.Series[int] + numpy array"""
58-
# i = np.array([2, 3, 5], np.int64)
59-
# f = np.array([1.0, 2.0, 3.0], np.float64)
60-
# c = np.array([1.1j, 2.2j, 4.1j], np.complex64)
59+
check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.int64)
60+
check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64)
61+
check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128)
6162

62-
# check(assert_type(left + i, "pd.Series[int]"), pd.Series, int)
63-
# check(assert_type(left + f, "pd.Series[float]"), pd.Series, float)
64-
# check(assert_type(left + c, "pd.Series[complex]"), pd.Series, complex)
63+
# check(assert_type(i + l, "pd.Series[int]"), pd.Series, np.int64)
64+
# check(assert_type(f + l, "pd.Series[float]"), pd.Series, np.float64)
65+
# check(assert_type(c + l, "pd.Series[complex]"), pd.Series, np.complex128)
6566

66-
# # check(assert_type(i + l, "pd.Series[int]"), pd.Series, int)
67-
# # check(assert_type(f + l, "pd.Series[float]"), pd.Series, float)
68-
# # check(assert_type(c + l, "pd.Series[complex]"), pd.Series, complex)
67+
check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.int64)
68+
check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64)
69+
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)
6970

70-
# check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, int)
71-
# check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, float)
72-
# check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, complex)
73-
74-
# check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, int)
75-
# check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, float)
76-
# check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, complex)
71+
check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.int64)
72+
check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64)
73+
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)
7774

7875

7976
def test_add_pd_series() -> None:
8077
"""Test pd.Series[int] + pandas series"""
8178
i = pd.Series([2, 3, 5])
82-
# f = pd.Series([1.0, 2.0, 3.0])
83-
# c = pd.Series([1.1j, 2.2j, 4.1j])
84-
85-
check(assert_type(left + i, "pd.Series[int]"), pd.Series, int)
86-
# check(assert_type(left + f, "pd.Series[float]"), pd.Series, float)
87-
# check(assert_type(left + c, "pd.Series[complex]"), pd.Series, complex)
88-
89-
# check(
90-
# assert_type(i + left, pd.Series), "pd.Series[int]"), pd.Series, int)
91-
# check(
92-
# assert_type(f + left, "pd.Series[float]"), pd.Series, float)
93-
# check(
94-
# assert_type(c + left, pd.Series), "pd.Series[complex]"), pd.Series, complex)
79+
f = pd.Series([1.0, 2.0, 3.0])
80+
c = pd.Series([1.1j, 2.2j, 4.1j])
9581

96-
check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, int)
97-
# check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, float)
98-
# check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, complex)
82+
check(assert_type(left + i, "pd.Series[int]"), pd.Series, np.int64)
83+
check(assert_type(left + f, "pd.Series[float]"), pd.Series, np.float64)
84+
check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128)
9985

100-
check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, int)
86+
check(assert_type(i + left, "pd.Series[int]"), pd.Series, np.int64)
87+
check(assert_type(f + left, "pd.Series[float]"), pd.Series, np.float64)
88+
check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128)
10189

90+
check(assert_type(left.add(i), "pd.Series[int]"), pd.Series, np.int64)
91+
check(assert_type(left.add(f), "pd.Series[float]"), pd.Series, np.float64)
92+
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)
10293

103-
# check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, float)
104-
# check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, complex)
94+
check(assert_type(left.radd(i), "pd.Series[int]"), pd.Series, np.int64)
95+
check(assert_type(left.radd(f), "pd.Series[float]"), pd.Series, np.float64)
96+
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)

‎tests/series/arithmetic/test_add.py

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -79,22 +79,13 @@ def test_add_pd_series() -> None:
7979
f = pd.Series([1.0, 2.0, 3.0])
8080
c = pd.Series([1.1j, 2.2j, 4.1j])
8181

82-
# check(assert_type(left + i, pd.Series), pd.Series) # why?
82+
check(assert_type(left + i, pd.Series), pd.Series)
8383
check(assert_type(left + f, pd.Series), pd.Series)
8484
check(assert_type(left + c, pd.Series), pd.Series)
8585

86-
check(
87-
assert_type(i + left, pd.Series), # pyright: ignore[reportAssertTypeFailure]
88-
pd.Series,
89-
)
90-
check(
91-
assert_type(f + left, pd.Series), # pyright: ignore[reportAssertTypeFailure]
92-
pd.Series,
93-
)
94-
check(
95-
assert_type(c + left, pd.Series), # pyright: ignore[reportAssertTypeFailure]
96-
pd.Series,
97-
)
86+
check(assert_type(i + left, pd.Series), pd.Series)
87+
check(assert_type(f + left, pd.Series), pd.Series)
88+
check(assert_type(c + left, pd.Series), pd.Series)
9889

9990
check(assert_type(left.add(i), pd.Series), pd.Series)
10091
check(assert_type(left.add(f), pd.Series), pd.Series)

‎tests/series/test_series.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1604,10 +1604,7 @@ def test_series_min_max_sub_axis() -> None:
16041604
ss = s1 - s2
16051605
sm = s1 * s2
16061606
sd = s1 / s2
1607-
check(
1608-
assert_type(sa, pd.Series), # pyright: ignore[reportAssertTypeFailure] # why?
1609-
pd.Series,
1610-
)
1607+
check(assert_type(sa, pd.Series), pd.Series)
16111608
check(assert_type(ss, pd.Series), pd.Series)
16121609
check(assert_type(sm, pd.Series), pd.Series)
16131610
check(assert_type(sd, pd.Series), pd.Series)

0 commit comments

Comments
 (0)