Skip to content

Commit fd4cedc

Browse files
cmp0xffcmp0xff
committed
feat: Series[complex] +
1 parent 8f43efc commit fd4cedc

File tree

3 files changed

+135
-1
lines changed

3 files changed

+135
-1
lines changed

‎pandas-stubs/core/series.pyi

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1620,7 +1620,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
16201620
@overload
16211621
def __add__(
16221622
self: Series[complex],
1623-
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1623+
other: (
1624+
_T_COMPLEX
1625+
| Sequence[_T_COMPLEX]
1626+
| Series[_T_COMPLEX]
1627+
| np_ndarray_anyint
1628+
| np_ndarray_float
1629+
| np_ndarray_complex
1630+
),
16241631
) -> Series[complex]: ...
16251632
@overload
16261633
def __add__(self, other: S1 | Self) -> Self: ...
@@ -1683,6 +1690,10 @@ class Series(IndexOpsMixin[S1], NDFrame):
16831690
self: Series[float], other: _T_COMPLEX | Sequence[_T_COMPLEX]
16841691
) -> Series[_T_COMPLEX]: ...
16851692
@overload
1693+
def __radd__(
1694+
self: Series[complex], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1695+
) -> Series[complex]: ...
1696+
@overload
16861697
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
16871698
@overload
16881699
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1832,6 +1843,20 @@ class Series(IndexOpsMixin[S1], NDFrame):
18321843
axis: int = ...,
18331844
) -> Series[complex]: ...
18341845
@overload
1846+
def add(
1847+
self: Series[complex],
1848+
other: (
1849+
Sequence[_T_COMPLEX]
1850+
| np_ndarray_anyint
1851+
| np_ndarray_float
1852+
| np_ndarray_complex
1853+
| Series[_T_COMPLEX]
1854+
),
1855+
level: Level | None = ...,
1856+
fill_value: float | None = ...,
1857+
axis: int = ...,
1858+
) -> Series[complex]: ...
1859+
@overload
18351860
def add(
18361861
self,
18371862
other: Series[S1] | Scalar,
@@ -2146,6 +2171,20 @@ class Series(IndexOpsMixin[S1], NDFrame):
21462171
axis: int = ...,
21472172
) -> Series[complex]: ...
21482173
@overload
2174+
def radd(
2175+
self: Series[complex],
2176+
other: (
2177+
Sequence[_T_COMPLEX]
2178+
| np_ndarray_anyint
2179+
| np_ndarray_float
2180+
| np_ndarray_complex
2181+
| Series[_T_COMPLEX]
2182+
),
2183+
level: Level | None = ...,
2184+
fill_value: float | None = ...,
2185+
axis: int = ...,
2186+
) -> Series[complex]: ...
2187+
@overload
21492188
def radd(
21502189
self,
21512190
other: Series[S1] | Scalar,

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

Whitespace-only changes.

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

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
import numpy as np
2+
import pandas as pd
3+
from typing_extensions import assert_type
4+
5+
from tests import check
6+
7+
left = pd.Series([1j, 2j, 3j]) # left operand
8+
9+
10+
def test_add_py_scalar() -> None:
11+
"""Test pd.Series[complex] + Python native scalars"""
12+
i, f, c = 1, 1.0, 1j
13+
14+
check(assert_type(left + i, "pd.Series[complex]"), pd.Series, np.complex128)
15+
check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complex128)
16+
check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128)
17+
18+
check(assert_type(i + left, "pd.Series[complex]"), pd.Series, np.complex128)
19+
check(assert_type(f + left, "pd.Series[complex]"), pd.Series, np.complex128)
20+
check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128)
21+
22+
check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128)
23+
check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128)
24+
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)
25+
26+
check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128)
27+
check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128)
28+
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)
29+
30+
31+
def test_add_py_sequence() -> None:
32+
"""Test pd.Series[complex] + Python native sequence"""
33+
i, f, c = [2, 3, 5], [1.0, 2.0, 3.0], [1j, 1j, 4j]
34+
35+
check(assert_type(left + i, "pd.Series[complex]"), pd.Series, np.complex128)
36+
check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complex128)
37+
check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128)
38+
39+
check(assert_type(i + left, "pd.Series[complex]"), pd.Series, np.complex128)
40+
check(assert_type(f + left, "pd.Series[complex]"), pd.Series, np.complex128)
41+
check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128)
42+
43+
check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128)
44+
check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128)
45+
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)
46+
47+
check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128)
48+
check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128)
49+
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)
50+
51+
52+
def test_add_numpy_array() -> None:
53+
"""Test pd.Series[complex] + numpy array"""
54+
i = np.array([2, 3, 5], np.int64)
55+
f = np.array([1.0, 2.0, 3.0], np.float64)
56+
c = np.array([1.1j, 2.2j, 4.1j], np.complex128)
57+
58+
check(assert_type(left + i, "pd.Series[complex]"), pd.Series, np.complex128)
59+
check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complex128)
60+
check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128)
61+
62+
# check(assert_type(i + l, "pd.Series[complex]"), pd.Series, np.complex128)
63+
# check(assert_type(f + l, "pd.Series[complex]"), pd.Series, np.complex128)
64+
# check(assert_type(c + l, "pd.Series[complex]"), pd.Series, np.complex128)
65+
66+
check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128)
67+
check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128)
68+
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)
69+
70+
check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128)
71+
check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128)
72+
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)
73+
74+
75+
def test_add_pd_series() -> None:
76+
"""Test pd.Series[complex] + pandas series"""
77+
i = pd.Series([2, 3, 5])
78+
f = pd.Series([1.0, 2.0, 3.0])
79+
c = pd.Series([1.1j, 2.2j, 4.1j])
80+
81+
check(assert_type(left + i, "pd.Series[complex]"), pd.Series, np.complex128)
82+
check(assert_type(left + f, "pd.Series[complex]"), pd.Series, np.complex128)
83+
check(assert_type(left + c, "pd.Series[complex]"), pd.Series, np.complex128)
84+
85+
check(assert_type(i + left, "pd.Series[complex]"), pd.Series, np.complex128)
86+
check(assert_type(f + left, "pd.Series[complex]"), pd.Series, np.complex128)
87+
check(assert_type(c + left, "pd.Series[complex]"), pd.Series, np.complex128)
88+
89+
check(assert_type(left.add(i), "pd.Series[complex]"), pd.Series, np.complex128)
90+
check(assert_type(left.add(f), "pd.Series[complex]"), pd.Series, np.complex128)
91+
check(assert_type(left.add(c), "pd.Series[complex]"), pd.Series, np.complex128)
92+
93+
check(assert_type(left.radd(i), "pd.Series[complex]"), pd.Series, np.complex128)
94+
check(assert_type(left.radd(f), "pd.Series[complex]"), pd.Series, np.complex128)
95+
check(assert_type(left.radd(c), "pd.Series[complex]"), pd.Series, np.complex128)

0 commit comments

Comments
 (0)