Skip to content

Commit 59af11e

Browse files
cmp0xffcmp0xff
committed
refactor: restart with simpler and systematic ideas
1 parent 197c0f9 commit 59af11e

File tree

5 files changed

+140
-167
lines changed

5 files changed

+140
-167
lines changed

‎pandas-stubs/core/series.pyi‎

Lines changed: 24 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ from typing import (
2525
Generic,
2626
Literal,
2727
NoReturn,
28-
TypeVar,
2928
overload,
3029
)
3130

@@ -175,6 +174,8 @@ from pandas._typing import (
175174
VoidDtypeArg,
176175
WriteBuffer,
177176
np_ndarray_anyint,
177+
np_ndarray_complex,
178+
np_ndarray_float,
178179
npt,
179180
num,
180181
)
@@ -184,12 +185,6 @@ from pandas.core.dtypes.dtypes import CategoricalDtype
184185

185186
from pandas.plotting import PlotAccessor
186187

187-
_T_INT = TypeVar("_T_INT", bound=int)
188-
_T_COMPLEX = TypeVar("_T_COMPLEX", bound=complex)
189-
_T_INT_FLOAT = TypeVar("_T_INT_FLOAT", bound=int | float)
190-
_T_FLOAT_COMPLEX = TypeVar("_T_FLOAT_COMPLEX", bound=float | complex)
191-
_T_INT_FLOAT_COMPLEX = TypeVar("_T_INT_FLOAT_COMPLEX", bound=int | float | complex)
192-
193188
class _iLocIndexerSeries(_iLocIndexer, Generic[S1]):
194189
# get item
195190
@overload
@@ -1588,22 +1583,8 @@ class Series(IndexOpsMixin[S1], NDFrame):
15881583
# them up.
15891584
@overload
15901585
def __add__(
1591-
self: Series[_T_INT], other: int | Sequence[int] | Series[int]
1592-
) -> Series[_T_INT]: ...
1593-
@overload
1594-
def __add__(
1595-
self: Series[_T_INT_FLOAT],
1596-
other: _T_FLOAT_COMPLEX | Sequence[_T_FLOAT_COMPLEX] | Series[_T_FLOAT_COMPLEX],
1597-
) -> Series[_T_FLOAT_COMPLEX]: ...
1598-
@overload
1599-
def __add__(
1600-
self: Series[_T_COMPLEX],
1601-
other: (
1602-
_T_INT_FLOAT_COMPLEX
1603-
| Sequence[_T_INT_FLOAT_COMPLEX]
1604-
| Series[_T_INT_FLOAT_COMPLEX]
1605-
),
1606-
) -> Series[_T_COMPLEX]: ...
1586+
self: Series[S1], other: complex | Sequence[complex] | Series[complex]
1587+
) -> Series[S1]: ...
16071588
@overload
16081589
def __add__(self, other: S1 | Self) -> Self: ...
16091590
@overload
@@ -1635,24 +1616,6 @@ class Series(IndexOpsMixin[S1], NDFrame):
16351616
self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date
16361617
) -> Series[_bool]: ...
16371618
@overload
1638-
def __mul__(
1639-
self: Series[_T_INT_FLOAT], other: int | Sequence[int] | Series[int]
1640-
) -> Series[_T_INT_FLOAT]: ...
1641-
@overload
1642-
def __mul__(
1643-
self: Series[_T_INT_FLOAT],
1644-
other: _T_FLOAT_COMPLEX | Sequence[_T_FLOAT_COMPLEX] | Series[_T_FLOAT_COMPLEX],
1645-
) -> Series[_T_FLOAT_COMPLEX]: ...
1646-
@overload
1647-
def __mul__(
1648-
self: Series[_T_COMPLEX],
1649-
other: (
1650-
_T_INT_FLOAT_COMPLEX
1651-
| Sequence[_T_INT_FLOAT_COMPLEX]
1652-
| Series[_T_INT_FLOAT_COMPLEX]
1653-
),
1654-
) -> Series[_T_COMPLEX]: ...
1655-
@overload
16561619
def __mul__(
16571620
self, other: timedelta | Timedelta | TimedeltaSeries | np.timedelta64
16581621
) -> TimedeltaSeries: ...
@@ -1669,6 +1632,10 @@ class Series(IndexOpsMixin[S1], NDFrame):
16691632
@overload
16701633
def __or__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
16711634
@overload
1635+
def __radd__(
1636+
self: Series[S1], other: complex | Sequence[complex] | np_ndarray_complex
1637+
) -> Series[S1]: ...
1638+
@overload
16721639
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
16731640
@overload
16741641
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
@@ -1746,26 +1713,16 @@ class Series(IndexOpsMixin[S1], NDFrame):
17461713
@property
17471714
def loc(self) -> _LocIndexerSeries[S1]: ...
17481715
# Methods
1749-
@overload
1750-
def add(
1751-
self: Series[_T_INT],
1752-
other: int | Sequence[int] | Series[int],
1753-
level: Level | None = ...,
1754-
fill_value: float | None = ...,
1755-
axis: int = ...,
1756-
) -> Series[_T_INT]: ...
1757-
@overload
1758-
def add(
1759-
self: Series[_T_INT_FLOAT],
1760-
other: _T_FLOAT_COMPLEX | Sequence[_T_FLOAT_COMPLEX] | Series[_T_FLOAT_COMPLEX],
1761-
level: Level | None = ...,
1762-
fill_value: float | None = ...,
1763-
axis: int = ...,
1764-
) -> Series[_T_FLOAT_COMPLEX]: ...
1765-
@overload
17661716
def add(
17671717
self,
1768-
other: Series[S1] | Scalar,
1718+
other: (
1719+
Sequence[S1]
1720+
| np_ndarray_anyint
1721+
| np_ndarray_float
1722+
| np_ndarray_complex
1723+
| Series[S1]
1724+
| Scalar
1725+
),
17691726
level: Level | None = ...,
17701727
fill_value: float | None = ...,
17711728
axis: int = ...,
@@ -2015,7 +1972,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
20151972
) -> Scalar: ...
20161973
def radd(
20171974
self,
2018-
other: Series[S1] | Scalar,
1975+
other: (
1976+
Sequence[S1]
1977+
| np_ndarray_anyint
1978+
| np_ndarray_float
1979+
| np_ndarray_complex
1980+
| Series[S1]
1981+
| Scalar
1982+
),
20191983
level: Level | None = ...,
20201984
fill_value: float | None = ...,
20211985
axis: AxisIndex = ...,

‎tests/series/__init__.py‎

Whitespace-only changes.

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

Whitespace-only changes.

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

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
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.DataFrame({"a": [1, 2, 3]})["a"] # left operand
9+
10+
11+
def test_add_py_scalar() -> None:
12+
"""Test pd.Series[Any] + Python native scalars"""
13+
i, f, c = 1, 1.0, 1j
14+
15+
check(assert_type(left + i, pd.Series), pd.Series)
16+
check(assert_type(left + f, pd.Series), pd.Series)
17+
check(assert_type(left + c, pd.Series), pd.Series)
18+
19+
check(assert_type(i + left, pd.Series), pd.Series)
20+
check(assert_type(f + left, pd.Series), pd.Series)
21+
check(assert_type(c + left, pd.Series), pd.Series)
22+
23+
check(assert_type(left.add(i), pd.Series), pd.Series)
24+
check(assert_type(left.add(f), pd.Series), pd.Series)
25+
check(assert_type(left.add(c), pd.Series), pd.Series)
26+
27+
check(assert_type(left.radd(i), pd.Series), pd.Series)
28+
check(assert_type(left.radd(f), pd.Series), pd.Series)
29+
check(assert_type(left.radd(c), pd.Series), pd.Series)
30+
31+
32+
def test_add_py_sequence() -> None:
33+
"""Test pd.Series[Any] + 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), pd.Series)
37+
check(assert_type(left + f, pd.Series), pd.Series)
38+
check(assert_type(left + c, pd.Series), pd.Series)
39+
40+
check(assert_type(i + left, pd.Series), pd.Series)
41+
check(assert_type(f + left, pd.Series), pd.Series)
42+
check(assert_type(c + left, pd.Series), pd.Series)
43+
44+
check(assert_type(left.add(i), pd.Series), pd.Series)
45+
check(assert_type(left.add(f), pd.Series), pd.Series)
46+
check(assert_type(left.add(c), pd.Series), pd.Series)
47+
48+
check(assert_type(left.radd(i), pd.Series), pd.Series)
49+
check(assert_type(left.radd(f), pd.Series), pd.Series)
50+
check(assert_type(left.radd(c), pd.Series), pd.Series)
51+
52+
53+
def test_add_numpy_array() -> None:
54+
"""Test pd.Series[Any] + 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.complex64)
58+
59+
check(assert_type(left + i, pd.Series), pd.Series)
60+
check(assert_type(left + f, pd.Series), pd.Series)
61+
check(assert_type(left + c, pd.Series), pd.Series)
62+
63+
# check(assert_type(i + l, pd.Series), pd.Series)
64+
# check(assert_type(f + l, pd.Series), pd.Series)
65+
# check(assert_type(c + l, pd.Series), pd.Series)
66+
67+
check(assert_type(left.add(i), pd.Series), pd.Series)
68+
check(assert_type(left.add(f), pd.Series), pd.Series)
69+
check(assert_type(left.add(c), pd.Series), pd.Series)
70+
71+
check(assert_type(left.radd(i), pd.Series), pd.Series)
72+
check(assert_type(left.radd(f), pd.Series), pd.Series)
73+
check(assert_type(left.radd(c), pd.Series), pd.Series)
74+
75+
76+
def test_add_pd_series() -> None:
77+
"""Test pd.Series[Any] + 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), pd.Series)
83+
check(assert_type(left + f, pd.Series), pd.Series)
84+
check(assert_type(left + c, pd.Series), pd.Series)
85+
86+
check(
87+
assert_type(i + left, pd.Series), # pyright: ignore[reportAssertTypeFailure]
88+
pd.Series,
89+
)
90+
check(
91+
assert_type(f + left, pd.Series), # pyright: ignore[reportAssertTypeFailure]
92+
pd.Series,
93+
)
94+
check(
95+
assert_type(c + left, pd.Series), # pyright: ignore[reportAssertTypeFailure]
96+
pd.Series,
97+
)
98+
99+
check(assert_type(left.add(i), pd.Series), pd.Series)
100+
check(assert_type(left.add(f), pd.Series), pd.Series)
101+
check(assert_type(left.add(c), pd.Series), pd.Series)
102+
103+
check(assert_type(left.radd(i), pd.Series), pd.Series)
104+
check(assert_type(left.radd(f), pd.Series), pd.Series)
105+
check(assert_type(left.radd(c), pd.Series), pd.Series)

0 commit comments

Comments
 (0)