Skip to content

Commit fcd829f

Browse files
loicdiridollouloicdiridollou
committed
GH1089 Partial typehinting
1 parent 3e56625 commit fcd829f

File tree

3 files changed

+327
-61
lines changed

3 files changed

+327
-61
lines changed

‎pandas-stubs/core/series.pyi

Lines changed: 133 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1609,6 +1609,16 @@ class Series(IndexOpsMixin[S1], NDFrame):
16091609
@overload
16101610
def __add__(self, other: S1 | Self) -> Self: ...
16111611
@overload
1612+
def __add__(
1613+
self: Series[int] | Series[float],
1614+
other: float | Series[float],
1615+
) -> Series[float]: ...
1616+
@overload
1617+
def __add__(
1618+
self: Series[float],
1619+
other: int | Series[int],
1620+
) -> Series[float]: ...
1621+
@overload
16121622
def __add__(
16131623
self,
16141624
other: complex,
@@ -1629,6 +1639,9 @@ class Series(IndexOpsMixin[S1], NDFrame):
16291639
) -> Series[int]: ...
16301640
# def __array__(self, dtype: Optional[_bool] = ...) -> _np_ndarray
16311641
def __eq__(self, other: object) -> Series[_bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
1642+
@overload
1643+
def __floordiv__(self: Series[float], other: Series[float] | Series[int] | float) -> Series[float]: ... # type: ignore[overload-overlap]
1644+
@overload
16321645
def __floordiv__(self, other: num | _ListLike | Series[S1]) -> Series[int]: ...
16331646
def __ge__( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
16341647
self, other: S1 | _ListLike | Series[S1] | datetime | timedelta | date
@@ -1657,11 +1670,29 @@ class Series(IndexOpsMixin[S1], NDFrame):
16571670
@overload
16581671
def __mul__(self: Series[int], other: Series[float] | float) -> Series[float]: ...
16591672
@overload
1673+
def __mul__(self: Series[float], other: Series[int] | int) -> Series[float]: ...
1674+
@overload
16601675
def __mul__(self: Series[Any], other: Series[Any]) -> Series: ...
16611676
@overload
16621677
def __mul__(self, other: num | _ListLike | Series) -> Series: ...
1678+
@overload
1679+
def __mod__( # pyright: ignore[reportOverlappingOverload]
1680+
self: Series[int], other: int | Series[int]
1681+
) -> Series[int]: ...
1682+
@overload
1683+
def __mod__(self: Series[int], other: float | Series[float]) -> Series[float]: ...
1684+
@overload
1685+
def __mod__(
1686+
self: Series[float], other: float | Series[int] | Series[float]
1687+
) -> Series[float]: ...
1688+
@overload
16631689
def __mod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
16641690
def __ne__(self, other: object) -> Series[_bool]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
1691+
@overload
1692+
def __pow__(self: Series[int], other: Series[int] | int) -> Series[int]: ...
1693+
@overload
1694+
def __pow__(self, other: Series[float] | float) -> Series[float]: ...
1695+
@overload
16651696
def __pow__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
16661697
# ignore needed for mypy as we want different results based on the arguments
16671698
@overload # type: ignore[override]
@@ -1686,7 +1717,12 @@ class Series(IndexOpsMixin[S1], NDFrame):
16861717
self, other: int | np_ndarray_anyint | Series[int]
16871718
) -> Series[int]: ...
16881719
def __rdivmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
1689-
def __rfloordiv__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
1720+
@overload
1721+
def __rfloordiv__(self: Series[float], other: float | Series[int] | Series[float]) -> Series[float]: ... # type: ignore[misc]
1722+
@overload
1723+
def __rfloordiv__(
1724+
self, other: num | _ListLike | Series[float]
1725+
) -> Series[float]: ...
16901726
def __rmod__(self, other: num | _ListLike | Series[S1]) -> Series[S1]: ...
16911727
@overload
16921728
def __rmul__(
@@ -1731,25 +1767,35 @@ class Series(IndexOpsMixin[S1], NDFrame):
17311767
self, other: Timestamp | datetime | TimestampSeries
17321768
) -> TimedeltaSeries: ...
17331769
@overload
1734-
def __sub__(
1770+
def __sub__(self, other: S1 | Self) -> Self: ...
1771+
@overload
1772+
def __sub__( # pyright: ignore[reportOverlappingOverload]
17351773
self: Series[int],
17361774
other: int | Series[int],
17371775
) -> Series[int]: ...
17381776
@overload
1777+
def __sub__(
1778+
self: Series[int],
1779+
other: float | Series[float],
1780+
) -> Series[float]: ...
1781+
@overload
1782+
def __sub__(
1783+
self: Series[float],
1784+
other: int | Series[int],
1785+
) -> Series[float]: ...
1786+
@overload
17391787
def __sub__(
17401788
self,
17411789
other: complex,
17421790
) -> Series[complex]: ...
17431791
@overload
1744-
def __sub__(self, other: S1 | Self) -> Self: ...
1745-
@overload
17461792
def __sub__(self, other: num | _ListLike | Series) -> Series: ...
17471793
@overload
17481794
def __truediv__(self: Series[int], other: Series[int] | int) -> Series[float]: ...
17491795
@overload
1750-
def __truediv__(
1751-
self, other: num | _ListLike | Series[S1] | Path
1752-
) -> Series | Self: ...
1796+
def __truediv__(self, other: Path) -> Series: ...
1797+
@overload
1798+
def __truediv__(self, other: num | _ListLike | Series[S1]) -> Series[float]: ...
17531799
# ignore needed for mypy as we want different results based on the arguments
17541800
@overload # type: ignore[override]
17551801
def __xor__( # pyright: ignore[reportOverlappingOverload]
@@ -1773,7 +1819,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
17731819
def iloc(self) -> _iLocIndexerSeries[S1]: ...
17741820
@property
17751821
def loc(self) -> _LocIndexerSeries[S1]: ...
1776-
# Methods
1822+
# Met @overload
17771823
@overload
17781824
def add(
17791825
self: Series[int],
@@ -1783,6 +1829,22 @@ class Series(IndexOpsMixin[S1], NDFrame):
17831829
axis: int = ...,
17841830
) -> Series[int]: ...
17851831
@overload
1832+
def add( # pyright: ignore[reportOverlappingOverload]
1833+
self,
1834+
other: float | Series[float],
1835+
level: Level | None = ...,
1836+
fill_value: float | None = ...,
1837+
axis: int = ...,
1838+
) -> Series[float]: ...
1839+
@overload
1840+
def add(
1841+
self: Series[float],
1842+
other: int | Series[int],
1843+
level: Level | None = ...,
1844+
fill_value: float | None = ...,
1845+
axis: int = ...,
1846+
) -> Series[float]: ...
1847+
@overload
17861848
def add(
17871849
self,
17881850
other: complex,
@@ -1798,6 +1860,14 @@ class Series(IndexOpsMixin[S1], NDFrame):
17981860
fill_value: float | None = ...,
17991861
axis: int = ...,
18001862
) -> Series[S1]: ...
1863+
@overload
1864+
def add(
1865+
self,
1866+
other: S1 | Self,
1867+
level: Level | None = ...,
1868+
fill_value: float | None = ...,
1869+
axis: int = ...,
1870+
) -> Series: ...
18011871
def all(
18021872
self,
18031873
axis: AxisIndex = ...,
@@ -1861,13 +1931,30 @@ class Series(IndexOpsMixin[S1], NDFrame):
18611931
min_periods: int = ...,
18621932
method: CalculationMethod = ...,
18631933
) -> Expanding[Series]: ...
1934+
@overload
1935+
def floordiv( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
1936+
self: Series[int],
1937+
other: Series[int] | int,
1938+
level: Level | None = ...,
1939+
fill_value: float | None = ...,
1940+
axis: AxisIndex | None = ...,
1941+
) -> Series[int]: ...
1942+
@overload
1943+
def floordiv(
1944+
self: Series[int] | Series[float],
1945+
other: num | _ListLike | Series[S1],
1946+
level: Level | None = ...,
1947+
fill_value: float | None = ...,
1948+
axis: AxisIndex | None = ...,
1949+
) -> Series[float]: ...
1950+
@overload
18641951
def floordiv(
18651952
self,
18661953
other: num | _ListLike | Series[S1],
18671954
level: Level | None = ...,
18681955
fill_value: float | None = ...,
18691956
axis: AxisIndex | None = ...,
1870-
) -> Series[int]: ...
1957+
) -> Series[S1]: ...
18711958
def ge(
18721959
self,
18731960
other: Scalar | Series[S1],
@@ -1945,6 +2032,31 @@ class Series(IndexOpsMixin[S1], NDFrame):
19452032
numeric_only: _bool = ...,
19462033
**kwargs,
19472034
) -> S1: ...
2035+
@overload
2036+
def mod( # pyright: ignore[reportOverlappingOverload]
2037+
self: Series[int],
2038+
other: Series[int] | int,
2039+
level: Level | None = ...,
2040+
fill_value: float | None = ...,
2041+
axis: AxisIndex | None = ...,
2042+
) -> Series[int]: ...
2043+
@overload
2044+
def mod(
2045+
self: Series[int],
2046+
other: Series[float] | float,
2047+
level: Level | None = ...,
2048+
fill_value: float | None = ...,
2049+
axis: AxisIndex | None = ...,
2050+
) -> Series[float]: ...
2051+
@overload
2052+
def mod(
2053+
self: Series[float],
2054+
other: Series[int] | Series[float] | float,
2055+
level: Level | None = ...,
2056+
fill_value: float | None = ...,
2057+
axis: AxisIndex | None = ...,
2058+
) -> Series[float]: ...
2059+
@overload
19482060
def mod(
19492061
self,
19502062
other: num | _ListLike | Series[S1],
@@ -1953,17 +2065,25 @@ class Series(IndexOpsMixin[S1], NDFrame):
19532065
axis: AxisIndex | None = ...,
19542066
) -> Series[S1]: ...
19552067
@overload
1956-
def mul(
2068+
def mul( # pyright: ignore[reportOverlappingOverload]
19572069
self: Series[int],
19582070
other: Series[int] | int,
19592071
level: Level | None = ...,
19602072
fill_value: int | None = ...,
19612073
axis: AxisIndex | None = ...,
19622074
) -> Series[int]: ...
19632075
@overload
2076+
def mul(
2077+
self: Series[float],
2078+
other: Series[int] | Series[float] | int,
2079+
level: Level | None = ...,
2080+
fill_value: float | None = ...,
2081+
axis: AxisIndex | None = ...,
2082+
) -> Series[float]: ...
2083+
@overload
19642084
def mul(
19652085
self: Series[int],
1966-
other: Series[float],
2086+
other: Series[float] | float,
19672087
level: Level | None = ...,
19682088
fill_value: float | None = ...,
19692089
axis: AxisIndex | None = ...,
@@ -2002,7 +2122,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
20022122
@overload
20032123
def pow(
20042124
self: Series[int],
2005-
other: float,
2125+
other: float | Series[float],
20062126
level: Level | None = ...,
20072127
fill_value: float | None = ...,
20082128
axis: AxisIndex | None = ...,
@@ -2161,7 +2281,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
21612281
@overload
21622282
def sub(
21632283
self: Series[int],
2164-
other: float,
2284+
other: float | Series[float],
21652285
level: Level | None = ...,
21662286
fill_value: float | None = ...,
21672287
axis: AxisIndex | None = ...,

‎tests/test_series.py

Lines changed: 3 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -646,45 +646,6 @@ def get_depth(url: str) -> int:
646646
check(assert_type(s.apply(lambda x: pd.NA), pd.Series), pd.Series, NAType)
647647

648648

649-
def test_types_element_wise_arithmetic() -> None:
650-
s = pd.Series([0, 1, -10])
651-
s2 = pd.Series([7, -5, 10])
652-
653-
check(assert_type(s + s2, "pd.Series[int]"), pd.Series, np.integer)
654-
check(assert_type(s.add(s2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer)
655-
656-
check(assert_type(s - s2, "pd.Series[int]"), pd.Series, np.integer)
657-
check(assert_type(s.sub(s2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer)
658-
659-
check(assert_type(s * s2, "pd.Series[int]"), pd.Series, np.integer)
660-
check(assert_type(s.mul(s2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer)
661-
662-
# GH1089 should be the following
663-
check(assert_type(s / s2, "pd.Series[float]"), pd.Series, np.float64)
664-
check(
665-
assert_type(s.div(s2, fill_value=0), "pd.Series[float]"), pd.Series, np.float64
666-
)
667-
668-
check(assert_type(s // s2, "pd.Series[int]"), pd.Series, np.integer)
669-
check(
670-
assert_type(s.floordiv(s2, fill_value=0), "pd.Series[int]"),
671-
pd.Series,
672-
np.integer,
673-
)
674-
675-
check(assert_type(s % s2, "pd.Series[int]"), pd.Series, np.integer)
676-
check(assert_type(s.mod(s2, fill_value=0), "pd.Series[int]"), pd.Series, np.integer)
677-
678-
check(assert_type(s ** s2.abs(), "pd.Series[int]"), pd.Series, np.integer)
679-
check(
680-
assert_type(s.pow(s2.abs(), fill_value=0), "pd.Series[int]"),
681-
pd.Series,
682-
np.integer,
683-
)
684-
685-
check(assert_type(divmod(s, s2), tuple["pd.Series[int]", "pd.Series[int]"]), tuple)
686-
687-
688649
def test_types_scalar_arithmetic() -> None:
689650
s = pd.Series([0, 1, -10])
690651

@@ -714,7 +675,7 @@ def test_types_scalar_arithmetic() -> None:
714675

715676
check(assert_type(s**2, "pd.Series[int]"), pd.Series, np.integer)
716677
check(assert_type(s**0, "pd.Series[int]"), pd.Series, np.integer)
717-
check(assert_type(s**0.213, "pd.Series[int]"), pd.Series, np.float64)
678+
check(assert_type(s**0.213, "pd.Series[float]"), pd.Series, np.float64)
718679
check(assert_type(s.pow(0.5), "pd.Series[float]"), pd.Series, np.float64)
719680

720681

@@ -1335,15 +1296,9 @@ def test_series_min_max_sub_axis() -> None:
13351296
sm = s1 * s2
13361297
sd = s1 / s2
13371298
check(assert_type(sa, pd.Series), pd.Series)
1338-
check(
1339-
assert_type(ss, pd.Series), # pyright: ignore[reportAssertTypeFailure]
1340-
pd.Series,
1341-
)
1299+
check(assert_type(ss, pd.Series), pd.Series)
13421300
check(assert_type(sm, pd.Series), pd.Series)
1343-
check(
1344-
assert_type(sd, pd.Series), # pyright: ignore[reportAssertTypeFailure]
1345-
pd.Series,
1346-
)
1301+
check(assert_type(sd, "pd.Series[float]"), pd.Series)
13471302

13481303

13491304
def test_series_index_isin() -> None:

0 commit comments

Comments
 (0)