Skip to content

Commit 566d8ec

Browse files
cmp0xffcmp0xff
committed
feat: Series[float] +
1 parent c494518 commit 566d8ec

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
@@ -1643,13 +1643,21 @@ class Series(IndexOpsMixin[S1], NDFrame):
16431643
@overload
16441644
def __add__(
16451645
self: Series[float],
1646-
other: int | Sequence[int] | np_ndarray_anyint | Series[int],
1646+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
16471647
) -> Series[float]: ...
16481648
@overload
1649+
def __add__(
1650+
self: Series[float],
1651+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1652+
) -> Series[_T_COMPLEX]: ...
1653+
@overload
1654+
def __add__(self: Series[float], other: np_ndarray_complex) -> Series[complex]: ...
1655+
@overload
16491656
def __add__(self: Series[complex], other: Series[Never]) -> Series: ...
16501657
@overload
16511658
def __add__(
1652-
self: Series[complex], other: int | Sequence[int] | Series[int]
1659+
self: Series[complex],
1660+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
16531661
) -> Series[complex]: ...
16541662
@overload
16551663
def __add__(self, other: S1 | Self) -> Self: ...
@@ -1706,6 +1714,12 @@ class Series(IndexOpsMixin[S1], NDFrame):
17061714
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX]
17071715
) -> Series[_T_COMPLEX]: ...
17081716
@overload
1717+
def __radd__(self: Series[float], other: int | Sequence[int]) -> Series[float]: ...
1718+
@overload
1719+
def __radd__(
1720+
self: Series[float], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1721+
) -> Series[_T_COMPLEX]: ...
1722+
@overload
17091723
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
17101724
@overload
17111725
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1832,6 +1846,30 @@ class Series(IndexOpsMixin[S1], NDFrame):
18321846
axis: int = ...,
18331847
) -> Series[complex]: ...
18341848
@overload
1849+
def add(
1850+
self: Series[float],
1851+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
1852+
level: Level | None = ...,
1853+
fill_value: float | None = ...,
1854+
axis: int = ...,
1855+
) -> Series[float]: ...
1856+
@overload
1857+
def add(
1858+
self: Series[float],
1859+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1860+
level: Level | None = ...,
1861+
fill_value: float | None = ...,
1862+
axis: int = ...,
1863+
) -> Series[_T_COMPLEX]: ...
1864+
@overload
1865+
def add(
1866+
self: Series[float],
1867+
other: np_ndarray_complex,
1868+
level: Level | None = ...,
1869+
fill_value: float | None = ...,
1870+
axis: int = ...,
1871+
) -> Series[complex]: ...
1872+
@overload
18351873
def add(
18361874
self,
18371875
other: Series[S1] | Scalar,
@@ -2124,6 +2162,30 @@ class Series(IndexOpsMixin[S1], NDFrame):
21242162
axis: int = ...,
21252163
) -> Series[complex]: ...
21262164
@overload
2165+
def radd(
2166+
self: Series[float],
2167+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
2168+
level: Level | None = ...,
2169+
fill_value: float | None = ...,
2170+
axis: int = ...,
2171+
) -> Series[float]: ...
2172+
@overload
2173+
def radd(
2174+
self: Series[float],
2175+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
2176+
level: Level | None = ...,
2177+
fill_value: float | None = ...,
2178+
axis: int = ...,
2179+
) -> Series[_T_COMPLEX]: ...
2180+
@overload
2181+
def radd(
2182+
self: Series[float],
2183+
other: np_ndarray_complex,
2184+
level: Level | None = ...,
2185+
fill_value: float | None = ...,
2186+
axis: int = ...,
2187+
) -> Series[complex]: ...
2188+
@overload
21272189
def radd(
21282190
self,
21292191
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)