Skip to content

Commit dbd9ea5

Browse files
cmp0xffcmp0xff
committed
feat: Series[int] +
1 parent 52d780c commit dbd9ea5

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
overload,
2930
)
3031

@@ -185,6 +186,8 @@ from pandas.core.dtypes.dtypes import CategoricalDtype
185186

186187
from pandas.plotting import PlotAccessor
187188

189+
_T_COMPLEX = TypeVar("_T_COMPLEX", bound=complex)
190+
188191
class _iLocIndexerSeries(_iLocIndexer, Generic[S1]):
189192
# get item
190193
@overload
@@ -1583,8 +1586,34 @@ class Series(IndexOpsMixin[S1], NDFrame):
15831586
# them up.
15841587
@overload
15851588
def __add__(
1586-
self: Series[int], other: int | Sequence[int] | Series[int]
1587-
) -> Series[int]: ...
1589+
self: Series[Never],
1590+
other: num | _str | timedelta | Timedelta | _ListLike | Series | np.timedelta64,
1591+
) -> Series: ...
1592+
@overload
1593+
def __add__(self: Series[int], other: Series[Never]) -> Series: ...
1594+
@overload
1595+
def __add__(
1596+
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX]
1597+
) -> Series[_T_COMPLEX]: ...
1598+
@overload
1599+
def __add__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ...
1600+
@overload
1601+
def __add__(self: Series[int], other: np_ndarray_float) -> Series[float]: ...
1602+
@overload
1603+
def __add__(self: Series[int], other: np_ndarray_complex) -> Series[complex]: ...
1604+
@overload
1605+
def __add__(self: Series[float], other: Series[Never]) -> Series: ...
1606+
@overload
1607+
def __add__(
1608+
self: Series[float],
1609+
other: int | Sequence[int] | np_ndarray_anyint | Series[int],
1610+
) -> Series[float]: ...
1611+
@overload
1612+
def __add__(self: Series[complex], other: Series[Never]) -> Series: ...
1613+
@overload
1614+
def __add__(
1615+
self: Series[complex], other: int | Sequence[int] | Series[int]
1616+
) -> Series[complex]: ...
15881617
@overload
15891618
def __add__(self, other: S1 | Self) -> Self: ...
15901619
@overload
@@ -1632,6 +1661,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
16321661
@overload
16331662
def __or__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
16341663
@overload
1664+
def __radd__(
1665+
self: Series[Never], other: num | _str | _ListLike | Series
1666+
) -> Series: ...
1667+
@overload
1668+
def __radd__(
1669+
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1670+
) -> Series[_T_COMPLEX]: ...
1671+
@overload
16351672
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
16361673
@overload
16371674
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1709,19 +1746,60 @@ class Series(IndexOpsMixin[S1], NDFrame):
17091746
@property
17101747
def loc(self) -> _LocIndexerSeries[S1]: ...
17111748
# Methods
1749+
@overload
17121750
def add(
1713-
self,
1751+
self: Series[Never],
17141752
other: (
1715-
Sequence[S1]
1753+
Sequence
17161754
| np_ndarray_anyint
17171755
| np_ndarray_float
17181756
| np_ndarray_complex
1719-
| Series[S1]
1757+
| Series
17201758
| Scalar
17211759
),
17221760
level: Level | None = ...,
17231761
fill_value: float | None = ...,
17241762
axis: int = ...,
1763+
) -> Series: ...
1764+
@overload
1765+
def add(
1766+
self: Series[int],
1767+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1768+
level: Level | None = ...,
1769+
fill_value: float | None = ...,
1770+
axis: int = ...,
1771+
) -> Series[_T_COMPLEX]: ...
1772+
@overload
1773+
def add(
1774+
self: Series[int],
1775+
other: np_ndarray_anyint,
1776+
level: Level | None = ...,
1777+
fill_value: float | None = ...,
1778+
axis: int = ...,
1779+
) -> Series[int]: ...
1780+
@overload
1781+
def add(
1782+
self: Series[int],
1783+
other: np_ndarray_float,
1784+
level: Level | None = ...,
1785+
fill_value: float | None = ...,
1786+
axis: int = ...,
1787+
) -> Series[float]: ...
1788+
@overload
1789+
def add(
1790+
self: Series[int],
1791+
other: np_ndarray_complex,
1792+
level: Level | None = ...,
1793+
fill_value: float | None = ...,
1794+
axis: int = ...,
1795+
) -> Series[complex]: ...
1796+
@overload
1797+
def add(
1798+
self,
1799+
other: Series[S1] | Scalar,
1800+
level: Level | None = ...,
1801+
fill_value: float | None = ...,
1802+
axis: int = ...,
17251803
) -> Series[S1]: ...
17261804
def all(
17271805
self,
@@ -1958,19 +2036,60 @@ class Series(IndexOpsMixin[S1], NDFrame):
19582036
min_count: int = ...,
19592037
**kwargs: Any,
19602038
) -> Scalar: ...
2039+
@overload
19612040
def radd(
1962-
self,
2041+
self: Series[Never],
19632042
other: (
1964-
Sequence[S1]
2043+
Sequence
19652044
| np_ndarray_anyint
19662045
| np_ndarray_float
19672046
| np_ndarray_complex
1968-
| Series[S1]
2047+
| Series
19692048
| Scalar
19702049
),
19712050
level: Level | None = ...,
19722051
fill_value: float | None = ...,
19732052
axis: AxisIndex = ...,
2053+
) -> Series: ...
2054+
@overload
2055+
def radd(
2056+
self: Series[int],
2057+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
2058+
level: Level | None = ...,
2059+
fill_value: float | None = ...,
2060+
axis: int = ...,
2061+
) -> Series[_T_COMPLEX]: ...
2062+
@overload
2063+
def radd(
2064+
self: Series[int],
2065+
other: np_ndarray_anyint,
2066+
level: Level | None = ...,
2067+
fill_value: float | None = ...,
2068+
axis: int = ...,
2069+
) -> Series[int]: ...
2070+
@overload
2071+
def radd(
2072+
self: Series[int],
2073+
other: np_ndarray_float,
2074+
level: Level | None = ...,
2075+
fill_value: float | None = ...,
2076+
axis: int = ...,
2077+
) -> Series[float]: ...
2078+
@overload
2079+
def radd(
2080+
self: Series[int],
2081+
other: np_ndarray_complex,
2082+
level: Level | None = ...,
2083+
fill_value: float | None = ...,
2084+
axis: int = ...,
2085+
) -> Series[complex]: ...
2086+
@overload
2087+
def radd(
2088+
self,
2089+
other: Series[S1] | Scalar,
2090+
level: Level | None = ...,
2091+
fill_value: float | None = ...,
2092+
axis: AxisIndex = ...,
19742093
) -> Series[S1]: ...
19752094
def rdivmod(
19762095
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)