Skip to content

Commit 46e904a

Browse files
cmp0xffcmp0xff
committed
feat: Series[float] +
1 parent cb7a17c commit 46e904a

File tree

3 files changed

+160
-2
lines changed

3 files changed

+160
-2
lines changed

‎pandas-stubs/core/series.pyi

Lines changed: 64 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1639,13 +1639,21 @@ class Series(IndexOpsMixin[S1], NDFrame):
16391639
@overload
16401640
def __add__(
16411641
self: Series[float],
1642-
other: int | Sequence[int] | np_ndarray_anyint | Series[int],
1642+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
16431643
) -> Series[float]: ...
16441644
@overload
1645+
def __add__(
1646+
self: Series[float],
1647+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1648+
) -> Series[_T_COMPLEX]: ...
1649+
@overload
1650+
def __add__(self: Series[float], other: np_ndarray_complex) -> Series[complex]: ...
1651+
@overload
16451652
def __add__(self: Series[complex], other: Series[Never]) -> Series: ...
16461653
@overload
16471654
def __add__(
1648-
self: Series[complex], other: int | Sequence[int] | Series[int]
1655+
self: Series[complex],
1656+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
16491657
) -> Series[complex]: ...
16501658
@overload
16511659
def __add__(self, other: S1 | Self) -> Self: ...
@@ -1702,6 +1710,12 @@ class Series(IndexOpsMixin[S1], NDFrame):
17021710
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX]
17031711
) -> Series[_T_COMPLEX]: ...
17041712
@overload
1713+
def __radd__(self: Series[float], other: int | Sequence[int]) -> Series[float]: ...
1714+
@overload
1715+
def __radd__(
1716+
self: Series[float], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1717+
) -> Series[_T_COMPLEX]: ...
1718+
@overload
17051719
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
17061720
@overload
17071721
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1828,6 +1842,30 @@ class Series(IndexOpsMixin[S1], NDFrame):
18281842
axis: int = ...,
18291843
) -> Series[complex]: ...
18301844
@overload
1845+
def add(
1846+
self: Series[float],
1847+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
1848+
level: Level | None = ...,
1849+
fill_value: float | None = ...,
1850+
axis: int = ...,
1851+
) -> Series[float]: ...
1852+
@overload
1853+
def add(
1854+
self: Series[float],
1855+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1856+
level: Level | None = ...,
1857+
fill_value: float | None = ...,
1858+
axis: int = ...,
1859+
) -> Series[_T_COMPLEX]: ...
1860+
@overload
1861+
def add(
1862+
self: Series[float],
1863+
other: np_ndarray_complex,
1864+
level: Level | None = ...,
1865+
fill_value: float | None = ...,
1866+
axis: int = ...,
1867+
) -> Series[complex]: ...
1868+
@overload
18311869
def add(
18321870
self,
18331871
other: Series[S1] | Scalar,
@@ -2122,6 +2160,30 @@ class Series(IndexOpsMixin[S1], NDFrame):
21222160
axis: int = ...,
21232161
) -> Series[complex]: ...
21242162
@overload
2163+
def radd(
2164+
self: Series[float],
2165+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
2166+
level: Level | None = ...,
2167+
fill_value: float | None = ...,
2168+
axis: int = ...,
2169+
) -> Series[float]: ...
2170+
@overload
2171+
def radd(
2172+
self: Series[float],
2173+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
2174+
level: Level | None = ...,
2175+
fill_value: float | None = ...,
2176+
axis: int = ...,
2177+
) -> Series[_T_COMPLEX]: ...
2178+
@overload
2179+
def radd(
2180+
self: Series[float],
2181+
other: np_ndarray_complex,
2182+
level: Level | None = ...,
2183+
fill_value: float | None = ...,
2184+
axis: int = ...,
2185+
) -> Series[complex]: ...
2186+
@overload
21252187
def radd(
21262188
self,
21272189
other: Series[S1] | Scalar,

‎tests/series/arithmetic/float/__init__.py

Whitespace-only changes.

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

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
from typing import assert_type
2+
3+
import numpy as np
4+
import pandas as pd
5+
6+
from tests import check
7+
8+
left = pd.Series([1.0, 2.0, 3.0]) # left operand
9+
10+
11+
def test_add_py_scalar() -> None:
12+
"""Test pd.Series[float] + Python native scalars"""
13+
i, f, c = 1, 1.0, 1j
14+
15+
check(assert_type(left + i, "pd.Series[float]"), pd.Series, np.float64)
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)
18+
19+
check(assert_type(i + left, "pd.Series[float]"), pd.Series, np.float64)
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)
22+
23+
check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.float64)
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)
26+
27+
check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.float64)
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)
30+
31+
32+
def test_add_py_sequence() -> None:
33+
"""Test pd.Series[float] + Python native sequence"""
34+
i, f, c = [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
35+
36+
check(assert_type(left + i, "pd.Series[float]"), pd.Series, np.float64)
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)
39+
40+
check(assert_type(i + left, "pd.Series[float]"), pd.Series, np.float64)
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)
43+
44+
check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.float64)
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)
47+
48+
check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.float64)
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)
51+
52+
53+
def test_add_numpy_array() -> None:
54+
"""Test pd.Series[float] + 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)
58+
59+
check(assert_type(left + i, "pd.Series[float]"), pd.Series, np.float64)
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)
62+
63+
# check(assert_type(i + l, "pd.Series[float]"), pd.Series, np.float64)
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)
66+
67+
check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.float64)
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)
70+
71+
check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.float64)
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)
74+
75+
76+
def test_add_pd_series() -> None:
77+
"""Test pd.Series[float] + pandas series"""
78+
i = pd.Series([2, 3, 5])
79+
f = pd.Series([1.0, 2.0, 3.0])
80+
c = pd.Series([1.1j, 2.2j, 4.1j])
81+
82+
check(assert_type(left + i, "pd.Series[float]"), pd.Series, np.float64)
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)
85+
86+
check(assert_type(i + left, "pd.Series[float]"), pd.Series, np.float64)
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)
89+
90+
check(assert_type(left.add(i), "pd.Series[float]"), pd.Series, np.float64)
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)
93+
94+
check(assert_type(left.radd(i), "pd.Series[float]"), pd.Series, np.float64)
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)

0 commit comments

Comments
 (0)