Skip to content

Commit

Permalink
GH1089 Partial typehinting
Browse files Browse the repository at this point in the history
  • Loading branch information
loicdiridollou committed Jan 15, 2025
1 parent 65aa064 commit 01130e1
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 6 deletions.
57 changes: 57 additions & 0 deletions pandas-stubs/core/series.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -1609,6 +1609,11 @@ class Series(IndexOpsMixin[S1], NDFrame):
@overload
def __add__(self, other: S1 | Self) -> Self: ...
@overload
def __add__(
self,
other: complex,
) -> Series[complex]: ...
@overload
def __add__(
self,
other: num | _str | timedelta | Timedelta | _ListLike | Series | np.timedelta64,
Expand Down Expand Up @@ -1716,6 +1721,16 @@ class Series(IndexOpsMixin[S1], NDFrame):
self, other: Timestamp | datetime | TimestampSeries
) -> TimedeltaSeries: ...
@overload
def __sub__(
self: Series[int],
other: int,
) -> Series[int]: ...
@overload
def __sub__(
self,
other: complex,
) -> Series[complex]: ...
@overload
def __sub__(self, other: num | _ListLike | Series) -> Series: ...
def __truediv__(self, other: num | _ListLike | Series[S1] | Path) -> Series: ...
# ignore needed for mypy as we want different results based on the arguments
Expand All @@ -1742,6 +1757,23 @@ class Series(IndexOpsMixin[S1], NDFrame):
@property
def loc(self) -> _LocIndexerSeries[S1]: ...
# Methods
@overload
def add(
self: Series[int],
other: int,
level: Level | None = ...,
fill_value: float | None = ...,
axis: int = ...,
) -> Series[int]: ...
@overload
def add(
self,
other: complex,
level: Level | None = ...,
fill_value: float | None = ...,
axis: int = ...,
) -> Series[complex]: ...
@overload
def add(
self,
other: Series[S1] | Scalar,
Expand Down Expand Up @@ -2085,6 +2117,31 @@ class Series(IndexOpsMixin[S1], NDFrame):
numeric_only: _bool = ...,
**kwargs,
) -> float: ...
@overload
def sub(
self: Series[int],
other: int,
level: Level | None = ...,
fill_value: float | None = ...,
axis: AxisIndex | None = ...,
) -> Series[int]: ...
@overload
def sub(
self: Series[int],
other: float,
level: Level | None = ...,
fill_value: float | None = ...,
axis: AxisIndex | None = ...,
) -> Series[int]: ...
@overload
def sub(
self,
other: complex,
level: Level | None = ...,
fill_value: float | None = ...,
axis: AxisIndex | None = ...,
) -> Series[complex]: ...
@overload
def sub(
self,
other: num | _ListLike | Series[S1],
Expand Down
14 changes: 8 additions & 6 deletions tests/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -450,7 +450,7 @@ def test_types_mean() -> None:
check(assert_type(s.mean(), float), float)
check(
assert_type(
s.groupby(level=0).mean(), # pyright: ignore[reportAssertTypeFailure]
s.groupby(level=0).mean(),
"pd.Series[float]",
),
pd.Series,
Expand All @@ -465,7 +465,7 @@ def test_types_median() -> None:
check(assert_type(s.median(), float), float)
check(
assert_type(
s.groupby(level=0).median(), # pyright: ignore[reportAssertTypeFailure]
s.groupby(level=0).median(),
"pd.Series[float]",
),
pd.Series,
Expand Down Expand Up @@ -690,7 +690,7 @@ def test_types_scalar_arithmetic() -> None:
check(assert_type(s + 1, "pd.Series[int]"), pd.Series, np.integer)
check(assert_type(s.add(1, fill_value=0), "pd.Series[int]"), pd.Series, np.integer)

check(assert_type(s - 1, pd.Series), pd.Series, np.integer)
check(assert_type(s - 1, "pd.Series[int]"), pd.Series, np.integer)
check(assert_type(s.sub(1, fill_value=0), "pd.Series[int]"), pd.Series, np.integer)

check(assert_type(s * 2, pd.Series), pd.Series, np.integer)
Expand Down Expand Up @@ -721,8 +721,10 @@ def test_types_scalar_arithmetic() -> None:
def test_types_complex_arithmetic() -> None:
c = 1 + 1j
s = pd.Series([1.0, 2.0, 3.0])
check(assert_type(s + c, pd.Series), pd.Series)
check(assert_type(s - c, pd.Series), pd.Series)
check(assert_type(s + c, "pd.Series[complex]"), pd.Series, complex)
check(assert_type(s.add(c), "pd.Series[complex]"), pd.Series, complex)
check(assert_type(s - c, "pd.Series[complex]"), pd.Series, complex)
check(assert_type(s.sub(c), "pd.Series[complex]"), pd.Series, complex)


def test_types_groupby() -> None:
Expand Down Expand Up @@ -1368,7 +1370,7 @@ def test_series_mul() -> None:
sm = s * 4
check(assert_type(sm, pd.Series), pd.Series)
ss = s - 4
check(assert_type(ss, pd.Series), pd.Series)
check(assert_type(ss, "pd.Series[int]"), pd.Series, np.integer)
sm2 = s * s
check(assert_type(sm2, pd.Series), pd.Series)
sp = s + 4
Expand Down

0 comments on commit 01130e1

Please sign in to comment.