Skip to content

Commit a63b9e4

Browse files
cmp0xffcmp0xff
committed
feat: Series[int] +
1 parent b3fda18 commit a63b9e4

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
@@ -1616,8 +1619,34 @@ class Series(IndexOpsMixin[S1], NDFrame):
16161619
# them up.
16171620
@overload
16181621
def __add__(
1619-
self: Series[int], other: int | Sequence[int] | Series[int]
1620-
) -> Series[int]: ...
1622+
self: Series[Never],
1623+
other: num | _str | timedelta | Timedelta | _ListLike | Series | np.timedelta64,
1624+
) -> Series: ...
1625+
@overload
1626+
def __add__(self: Series[int], other: Series[Never]) -> Series: ...
1627+
@overload
1628+
def __add__(
1629+
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX]
1630+
) -> Series[_T_COMPLEX]: ...
1631+
@overload
1632+
def __add__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ...
1633+
@overload
1634+
def __add__(self: Series[int], other: np_ndarray_float) -> Series[float]: ...
1635+
@overload
1636+
def __add__(self: Series[int], other: np_ndarray_complex) -> Series[complex]: ...
1637+
@overload
1638+
def __add__(self: Series[float], other: Series[Never]) -> Series: ...
1639+
@overload
1640+
def __add__(
1641+
self: Series[float],
1642+
other: int | Sequence[int] | np_ndarray_anyint | Series[int],
1643+
) -> Series[float]: ...
1644+
@overload
1645+
def __add__(self: Series[complex], other: Series[Never]) -> Series: ...
1646+
@overload
1647+
def __add__(
1648+
self: Series[complex], other: int | Sequence[int] | Series[int]
1649+
) -> Series[complex]: ...
16211650
@overload
16221651
def __add__(self, other: S1 | Self) -> Self: ...
16231652
@overload
@@ -1665,6 +1694,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
16651694
@overload
16661695
def __or__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
16671696
@overload
1697+
def __radd__(
1698+
self: Series[Never], other: num | _str | _ListLike | Series
1699+
) -> Series: ...
1700+
@overload
1701+
def __radd__(
1702+
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1703+
) -> Series[_T_COMPLEX]: ...
1704+
@overload
16681705
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
16691706
@overload
16701707
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1743,19 +1780,60 @@ class Series(IndexOpsMixin[S1], NDFrame):
17431780
@property
17441781
def loc(self) -> _LocIndexerSeries[S1]: ...
17451782
# Methods
1783+
@overload
17461784
def add(
1747-
self,
1785+
self: Series[Never],
17481786
other: (
1749-
Sequence[S1]
1787+
Sequence
17501788
| np_ndarray_anyint
17511789
| np_ndarray_float
17521790
| np_ndarray_complex
1753-
| Series[S1]
1791+
| Series
17541792
| Scalar
17551793
),
17561794
level: Level | None = ...,
17571795
fill_value: float | None = ...,
17581796
axis: int = ...,
1797+
) -> Series: ...
1798+
@overload
1799+
def add(
1800+
self: Series[int],
1801+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1802+
level: Level | None = ...,
1803+
fill_value: float | None = ...,
1804+
axis: int = ...,
1805+
) -> Series[_T_COMPLEX]: ...
1806+
@overload
1807+
def add(
1808+
self: Series[int],
1809+
other: np_ndarray_anyint,
1810+
level: Level | None = ...,
1811+
fill_value: float | None = ...,
1812+
axis: int = ...,
1813+
) -> Series[int]: ...
1814+
@overload
1815+
def add(
1816+
self: Series[int],
1817+
other: np_ndarray_float,
1818+
level: Level | None = ...,
1819+
fill_value: float | None = ...,
1820+
axis: int = ...,
1821+
) -> Series[float]: ...
1822+
@overload
1823+
def add(
1824+
self: Series[int],
1825+
other: np_ndarray_complex,
1826+
level: Level | None = ...,
1827+
fill_value: float | None = ...,
1828+
axis: int = ...,
1829+
) -> Series[complex]: ...
1830+
@overload
1831+
def add(
1832+
self,
1833+
other: Series[S1] | Scalar,
1834+
level: Level | None = ...,
1835+
fill_value: float | None = ...,
1836+
axis: int = ...,
17591837
) -> Series[S1]: ...
17601838
def all(
17611839
self,
@@ -1996,19 +2074,60 @@ class Series(IndexOpsMixin[S1], NDFrame):
19962074
min_count: int = ...,
19972075
**kwargs: Any,
19982076
) -> Scalar: ...
2077+
@overload
19992078
def radd(
2000-
self,
2079+
self: Series[Never],
20012080
other: (
2002-
Sequence[S1]
2081+
Sequence
20032082
| np_ndarray_anyint
20042083
| np_ndarray_float
20052084
| np_ndarray_complex
2006-
| Series[S1]
2085+
| Series
20072086
| Scalar
20082087
),
20092088
level: Level | None = ...,
20102089
fill_value: float | None = ...,
20112090
axis: AxisIndex = ...,
2091+
) -> Series: ...
2092+
@overload
2093+
def radd(
2094+
self: Series[int],
2095+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
2096+
level: Level | None = ...,
2097+
fill_value: float | None = ...,
2098+
axis: int = ...,
2099+
) -> Series[_T_COMPLEX]: ...
2100+
@overload
2101+
def radd(
2102+
self: Series[int],
2103+
other: np_ndarray_anyint,
2104+
level: Level | None = ...,
2105+
fill_value: float | None = ...,
2106+
axis: int = ...,
2107+
) -> Series[int]: ...
2108+
@overload
2109+
def radd(
2110+
self: Series[int],
2111+
other: np_ndarray_float,
2112+
level: Level | None = ...,
2113+
fill_value: float | None = ...,
2114+
axis: int = ...,
2115+
) -> Series[float]: ...
2116+
@overload
2117+
def radd(
2118+
self: Series[int],
2119+
other: np_ndarray_complex,
2120+
level: Level | None = ...,
2121+
fill_value: float | None = ...,
2122+
axis: int = ...,
2123+
) -> Series[complex]: ...
2124+
@overload
2125+
def radd(
2126+
self,
2127+
other: Series[S1] | Scalar,
2128+
level: Level | None = ...,
2129+
fill_value: float | None = ...,
2130+
axis: AxisIndex = ...,
20122131
) -> Series[S1]: ...
20132132
def rdivmod(
20142133
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
@@ -1599,10 +1599,7 @@ def test_series_min_max_sub_axis() -> None:
15991599
ss = s1 - s2
16001600
sm = s1 * s2
16011601
sd = s1 / s2
1602-
check(
1603-
assert_type(sa, pd.Series), # pyright: ignore[reportAssertTypeFailure] # why?
1604-
pd.Series,
1605-
)
1602+
check(assert_type(sa, pd.Series), pd.Series)
16061603
check(assert_type(ss, pd.Series), pd.Series)
16071604
check(assert_type(sm, pd.Series), pd.Series)
16081605
check(assert_type(sd, pd.Series), pd.Series)

0 commit comments

Comments
 (0)