Skip to content

Commit 6f8e00a

Browse files
cmp0xffcmp0xff
committed
feat: Series[complex] +
1 parent 2581638 commit 6f8e00a

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
@@ -1657,7 +1657,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
16571657
@overload
16581658
def __add__(
16591659
self: Series[complex],
1660-
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1660+
other: (
1661+
_T_COMPLEX
1662+
| Sequence[_T_COMPLEX]
1663+
| Series[_T_COMPLEX]
1664+
| np_ndarray_anyint
1665+
| np_ndarray_float
1666+
| np_ndarray_complex
1667+
),
16611668
) -> Series[complex]: ...
16621669
@overload
16631670
def __add__(self, other: S1 | Self) -> Self: ...
@@ -1720,6 +1727,10 @@ class Series(IndexOpsMixin[S1], NDFrame):
17201727
self: Series[float], other: _T_COMPLEX | Sequence[_T_COMPLEX]
17211728
) -> Series[_T_COMPLEX]: ...
17221729
@overload
1730+
def __radd__(
1731+
self: Series[complex], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1732+
) -> Series[complex]: ...
1733+
@overload
17231734
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
17241735
@overload
17251736
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1870,6 +1881,20 @@ class Series(IndexOpsMixin[S1], NDFrame):
18701881
axis: int = ...,
18711882
) -> Series[complex]: ...
18721883
@overload
1884+
def add(
1885+
self: Series[complex],
1886+
other: (
1887+
Sequence[_T_COMPLEX]
1888+
| np_ndarray_anyint
1889+
| np_ndarray_float
1890+
| np_ndarray_complex
1891+
| Series[_T_COMPLEX]
1892+
),
1893+
level: Level | None = ...,
1894+
fill_value: float | None = ...,
1895+
axis: int = ...,
1896+
) -> Series[complex]: ...
1897+
@overload
18731898
def add(
18741899
self,
18751900
other: Series[S1] | Scalar,
@@ -2186,6 +2211,20 @@ class Series(IndexOpsMixin[S1], NDFrame):
21862211
axis: int = ...,
21872212
) -> Series[complex]: ...
21882213
@overload
2214+
def radd(
2215+
self: Series[complex],
2216+
other: (
2217+
Sequence[_T_COMPLEX]
2218+
| np_ndarray_anyint
2219+
| np_ndarray_float
2220+
| np_ndarray_complex
2221+
| Series[_T_COMPLEX]
2222+
),
2223+
level: Level | None = ...,
2224+
fill_value: float | None = ...,
2225+
axis: int = ...,
2226+
) -> Series[complex]: ...
2227+
@overload
21892228
def radd(
21902229
self,
21912230
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)