Skip to content

Commit e3bfc89

Browse files
cmp0xffcmp0xff
committed
feat: Series[complex] +
1 parent 2886a29 commit e3bfc89

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
@@ -1653,7 +1653,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
16531653
@overload
16541654
def __add__(
16551655
self: Series[complex],
1656-
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1656+
other: (
1657+
_T_COMPLEX
1658+
| Sequence[_T_COMPLEX]
1659+
| Series[_T_COMPLEX]
1660+
| np_ndarray_anyint
1661+
| np_ndarray_float
1662+
| np_ndarray_complex
1663+
),
16571664
) -> Series[complex]: ...
16581665
@overload
16591666
def __add__(self, other: S1 | Self) -> Self: ...
@@ -1716,6 +1723,10 @@ class Series(IndexOpsMixin[S1], NDFrame):
17161723
self: Series[float], other: _T_COMPLEX | Sequence[_T_COMPLEX]
17171724
) -> Series[_T_COMPLEX]: ...
17181725
@overload
1726+
def __radd__(
1727+
self: Series[complex], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1728+
) -> Series[complex]: ...
1729+
@overload
17191730
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
17201731
@overload
17211732
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1866,6 +1877,20 @@ class Series(IndexOpsMixin[S1], NDFrame):
18661877
axis: int = ...,
18671878
) -> Series[complex]: ...
18681879
@overload
1880+
def add(
1881+
self: Series[complex],
1882+
other: (
1883+
Sequence[_T_COMPLEX]
1884+
| np_ndarray_anyint
1885+
| np_ndarray_float
1886+
| np_ndarray_complex
1887+
| Series[_T_COMPLEX]
1888+
),
1889+
level: Level | None = ...,
1890+
fill_value: float | None = ...,
1891+
axis: int = ...,
1892+
) -> Series[complex]: ...
1893+
@overload
18691894
def add(
18701895
self,
18711896
other: Series[S1] | Scalar,
@@ -2184,6 +2209,20 @@ class Series(IndexOpsMixin[S1], NDFrame):
21842209
axis: int = ...,
21852210
) -> Series[complex]: ...
21862211
@overload
2212+
def radd(
2213+
self: Series[complex],
2214+
other: (
2215+
Sequence[_T_COMPLEX]
2216+
| np_ndarray_anyint
2217+
| np_ndarray_float
2218+
| np_ndarray_complex
2219+
| Series[_T_COMPLEX]
2220+
),
2221+
level: Level | None = ...,
2222+
fill_value: float | None = ...,
2223+
axis: int = ...,
2224+
) -> Series[complex]: ...
2225+
@overload
21872226
def radd(
21882227
self,
21892228
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)