Skip to content

Commit 845f9c5

Browse files
cmp0xffcmp0xff
authored
feat(series): #1098 arithmetic addition (#1275)
* refactor: minimal systematic implementation * feat: Series[int] + * feat: Series[float] + * refactor: py310 and remove why * feat: Series[complex] + * fix(comment): #1275 (comment) * fix(comment): https://github.com/pandas-dev/pandas-stubs/pull/1275/files#r2211559024 * fix(comment): https://github.com/pandas-dev/pandas-stubs/pull/1275/files#r2211571551 * fix(comment): argument types #1275 (review) * fix(comment): #1275 (review)
1 parent 21dfb77 commit 845f9c5

File tree

12 files changed

+650
-12
lines changed

12 files changed

+650
-12
lines changed

‎pandas-stubs/core/series.pyi

Lines changed: 241 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ from typing import (
2525
Generic,
2626
Literal,
2727
NoReturn,
28+
TypeVar,
2829
final,
2930
overload,
3031
)
@@ -175,6 +176,8 @@ from pandas._typing import (
175176
VoidDtypeArg,
176177
WriteBuffer,
177178
np_ndarray_anyint,
179+
np_ndarray_complex,
180+
np_ndarray_float,
178181
npt,
179182
num,
180183
)
@@ -184,6 +187,8 @@ from pandas.core.dtypes.dtypes import CategoricalDtype
184187

185188
from pandas.plotting import PlotAccessor
186189

190+
_T_COMPLEX = TypeVar("_T_COMPLEX", bound=complex)
191+
187192
class _iLocIndexerSeries(_iLocIndexer, Generic[S1]):
188193
# get item
189194
@overload
@@ -1617,12 +1622,48 @@ class Series(IndexOpsMixin[S1], NDFrame):
16171622
# just failed to generate these so I couldn't match
16181623
# them up.
16191624
@overload
1620-
def __add__(self, other: S1 | Self) -> Self: ...
1621-
@overload
16221625
def __add__(
1623-
self,
1624-
other: num | _str | timedelta | Timedelta | _ListLike | Series | np.timedelta64,
1626+
self: Series[Never],
1627+
other: Scalar | _ListLike | Series,
16251628
) -> Series: ...
1629+
@overload
1630+
def __add__(self, other: Series[Never]) -> Series: ...
1631+
@overload
1632+
def __add__(
1633+
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX]
1634+
) -> Series[_T_COMPLEX]: ...
1635+
@overload
1636+
def __add__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ...
1637+
@overload
1638+
def __add__(self: Series[int], other: np_ndarray_float) -> Series[float]: ...
1639+
@overload
1640+
def __add__(self: Series[int], other: np_ndarray_complex) -> Series[complex]: ...
1641+
@overload
1642+
def __add__(
1643+
self: Series[float],
1644+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
1645+
) -> Series[float]: ...
1646+
@overload
1647+
def __add__(
1648+
self: Series[float],
1649+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1650+
) -> Series[_T_COMPLEX]: ...
1651+
@overload
1652+
def __add__(self: Series[float], other: np_ndarray_complex) -> Series[complex]: ...
1653+
@overload
1654+
def __add__(
1655+
self: Series[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+
),
1664+
) -> Series[complex]: ...
1665+
@overload
1666+
def __add__(self, other: S1 | Self) -> Self: ...
16261667
# ignore needed for mypy as we want different results based on the arguments
16271668
@overload # type: ignore[override]
16281669
def __and__( # pyright: ignore[reportOverlappingOverload]
@@ -1663,9 +1704,40 @@ class Series(IndexOpsMixin[S1], NDFrame):
16631704
@overload
16641705
def __or__(self, other: int | np_ndarray_anyint | Series[int]) -> Series[int]: ...
16651706
@overload
1666-
def __radd__(self, other: S1 | Series[S1]) -> Self: ...
1707+
def __radd__(self: Series[Never], other: Scalar | _ListLike) -> Series: ...
1708+
@overload
1709+
def __radd__(
1710+
self: Series[int], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1711+
) -> Series[_T_COMPLEX]: ...
1712+
@overload
1713+
def __radd__(self: Series[int], other: np_ndarray_anyint) -> Series[int]: ...
1714+
@overload
1715+
def __radd__(self: Series[int], other: np_ndarray_float) -> Series[float]: ...
1716+
@overload
1717+
def __radd__(self: Series[int], other: np_ndarray_complex) -> Series[complex]: ...
1718+
@overload
1719+
def __radd__(
1720+
self: Series[float],
1721+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float,
1722+
) -> Series[float]: ...
1723+
@overload
1724+
def __radd__(
1725+
self: Series[float], other: _T_COMPLEX | Sequence[_T_COMPLEX]
1726+
) -> Series[_T_COMPLEX]: ...
16671727
@overload
1668-
def __radd__(self, other: num | _str | _ListLike | Series) -> Series: ...
1728+
def __radd__(self: Series[float], other: np_ndarray_complex) -> Series[complex]: ...
1729+
@overload
1730+
def __radd__(
1731+
self: Series[complex],
1732+
other: (
1733+
np_ndarray_anyint
1734+
| np_ndarray_float
1735+
| np_ndarray_complex
1736+
| Sequence[_T_COMPLEX]
1737+
),
1738+
) -> Series[complex]: ...
1739+
@overload
1740+
def __radd__(self, other: S1) -> Self: ...
16691741
# ignore needed for mypy as we want different results based on the arguments
16701742
@overload # type: ignore[override]
16711743
def __rand__( # pyright: ignore[reportOverlappingOverload]
@@ -1741,13 +1813,92 @@ class Series(IndexOpsMixin[S1], NDFrame):
17411813
@property
17421814
def loc(self) -> _LocIndexerSeries[S1]: ...
17431815
# Methods
1816+
@overload
1817+
def add(
1818+
self: Series[Never],
1819+
other: Scalar | _ListLike | Series,
1820+
level: Level | None = ...,
1821+
fill_value: float | None = ...,
1822+
axis: int = ...,
1823+
) -> Series: ...
1824+
@overload
1825+
def add(
1826+
self: Series[int],
1827+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1828+
level: Level | None = ...,
1829+
fill_value: float | None = ...,
1830+
axis: int = ...,
1831+
) -> Series[_T_COMPLEX]: ...
1832+
@overload
1833+
def add(
1834+
self: Series[int],
1835+
other: np_ndarray_anyint,
1836+
level: Level | None = ...,
1837+
fill_value: float | None = ...,
1838+
axis: int = ...,
1839+
) -> Series[int]: ...
1840+
@overload
1841+
def add(
1842+
self: Series[int],
1843+
other: np_ndarray_float,
1844+
level: Level | None = ...,
1845+
fill_value: float | None = ...,
1846+
axis: int = ...,
1847+
) -> Series[float]: ...
1848+
@overload
1849+
def add(
1850+
self: Series[int],
1851+
other: np_ndarray_complex,
1852+
level: Level | None = ...,
1853+
fill_value: float | None = ...,
1854+
axis: int = ...,
1855+
) -> Series[complex]: ...
1856+
@overload
1857+
def add(
1858+
self: Series[float],
1859+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
1860+
level: Level | None = ...,
1861+
fill_value: float | None = ...,
1862+
axis: int = ...,
1863+
) -> Series[float]: ...
1864+
@overload
1865+
def add(
1866+
self: Series[float],
1867+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
1868+
level: Level | None = ...,
1869+
fill_value: float | None = ...,
1870+
axis: int = ...,
1871+
) -> Series[_T_COMPLEX]: ...
1872+
@overload
1873+
def add(
1874+
self: Series[float],
1875+
other: np_ndarray_complex,
1876+
level: Level | None = ...,
1877+
fill_value: float | None = ...,
1878+
axis: int = ...,
1879+
) -> Series[complex]: ...
1880+
@overload
1881+
def add(
1882+
self: Series[complex],
1883+
other: (
1884+
Sequence[_T_COMPLEX]
1885+
| np_ndarray_anyint
1886+
| np_ndarray_float
1887+
| np_ndarray_complex
1888+
| Series[_T_COMPLEX]
1889+
),
1890+
level: Level | None = ...,
1891+
fill_value: float | None = ...,
1892+
axis: int = ...,
1893+
) -> Series[complex]: ...
1894+
@overload
17441895
def add(
17451896
self,
1746-
other: Series[S1] | Scalar,
1897+
other: S1 | Series[S1],
17471898
level: Level | None = ...,
17481899
fill_value: float | None = ...,
17491900
axis: int = ...,
1750-
) -> Series[S1]: ...
1901+
) -> Self: ...
17511902
def all(
17521903
self,
17531904
axis: AxisIndex = ...,
@@ -1985,13 +2136,92 @@ class Series(IndexOpsMixin[S1], NDFrame):
19852136
min_count: int = ...,
19862137
**kwargs: Any,
19872138
) -> Scalar: ...
2139+
@overload
2140+
def radd(
2141+
self: Series[Never],
2142+
other: Scalar | _ListLike | Series,
2143+
level: Level | None = ...,
2144+
fill_value: float | None = ...,
2145+
axis: AxisIndex = ...,
2146+
) -> Series: ...
2147+
@overload
2148+
def radd(
2149+
self: Series[int],
2150+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
2151+
level: Level | None = ...,
2152+
fill_value: float | None = ...,
2153+
axis: int = ...,
2154+
) -> Series[_T_COMPLEX]: ...
2155+
@overload
2156+
def radd(
2157+
self: Series[int],
2158+
other: np_ndarray_anyint,
2159+
level: Level | None = ...,
2160+
fill_value: float | None = ...,
2161+
axis: int = ...,
2162+
) -> Series[int]: ...
2163+
@overload
2164+
def radd(
2165+
self: Series[int],
2166+
other: np_ndarray_float,
2167+
level: Level | None = ...,
2168+
fill_value: float | None = ...,
2169+
axis: int = ...,
2170+
) -> Series[float]: ...
2171+
@overload
2172+
def radd(
2173+
self: Series[int],
2174+
other: np_ndarray_complex,
2175+
level: Level | None = ...,
2176+
fill_value: float | None = ...,
2177+
axis: int = ...,
2178+
) -> Series[complex]: ...
2179+
@overload
2180+
def radd(
2181+
self: Series[float],
2182+
other: int | Sequence[int] | np_ndarray_anyint | np_ndarray_float | Series[int],
2183+
level: Level | None = ...,
2184+
fill_value: float | None = ...,
2185+
axis: int = ...,
2186+
) -> Series[float]: ...
2187+
@overload
2188+
def radd(
2189+
self: Series[float],
2190+
other: _T_COMPLEX | Sequence[_T_COMPLEX] | Series[_T_COMPLEX],
2191+
level: Level | None = ...,
2192+
fill_value: float | None = ...,
2193+
axis: int = ...,
2194+
) -> Series[_T_COMPLEX]: ...
2195+
@overload
2196+
def radd(
2197+
self: Series[float],
2198+
other: np_ndarray_complex,
2199+
level: Level | None = ...,
2200+
fill_value: float | None = ...,
2201+
axis: int = ...,
2202+
) -> Series[complex]: ...
2203+
@overload
2204+
def radd(
2205+
self: Series[complex],
2206+
other: (
2207+
Sequence[_T_COMPLEX]
2208+
| np_ndarray_anyint
2209+
| np_ndarray_float
2210+
| np_ndarray_complex
2211+
| Series[_T_COMPLEX]
2212+
),
2213+
level: Level | None = ...,
2214+
fill_value: float | None = ...,
2215+
axis: int = ...,
2216+
) -> Series[complex]: ...
2217+
@overload
19882218
def radd(
19892219
self,
1990-
other: Series[S1] | Scalar,
2220+
other: S1 | Series[S1],
19912221
level: Level | None = ...,
19922222
fill_value: float | None = ...,
19932223
axis: AxisIndex = ...,
1994-
) -> Series[S1]: ...
2224+
) -> Self: ...
19952225
def rdivmod(
19962226
self,
19972227
other: Series[S1] | Scalar,
@@ -2395,7 +2625,7 @@ class PeriodSeries(Series[Period]):
23952625
) -> Never: ...
23962626

23972627
class OffsetSeries(Series[BaseOffset]):
2398-
@overload # type: ignore[override]
2628+
@overload
23992629
def __radd__(self, other: Period) -> PeriodSeries: ...
24002630
@overload
24012631
def __radd__( # pyright: ignore[reportIncompatibleMethodOverride]

‎pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ mypy = "1.17.0"
3939
pandas = "2.3.1"
4040
pyarrow = ">=10.0.1"
4141
pytest = ">=7.1.2"
42-
pyright = ">=1.1.400"
42+
pyright = ">=1.1.403"
4343
ty = "^0.0.1a8"
4444
pyrefly = "^0.21.0"
4545
poethepoet = ">=0.16.5"

‎tests/series/__init__.py

Whitespace-only changes.

‎tests/series/arithmetic/__init__.py

Whitespace-only changes.

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

Whitespace-only changes.

0 commit comments

Comments
 (0)