diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index cbaa3096..165f1614 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -67,6 +67,7 @@ from pandas._libs.lib import NoDefault from pandas._libs.missing import NAType from pandas._libs.tslibs import BaseOffset from pandas._libs.tslibs.nattype import NaTType +from pandas._libs.tslibs.offsets import DateOffset from pandas._typing import ( S1, AggFuncTypeBase, @@ -87,7 +88,6 @@ from pandas._typing import ( FilePath, FillnaOptions, FormattersType, - Frequency, GroupByObjectNonScalar, HashableT, HashableT1, @@ -830,10 +830,10 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): ) -> Self: ... def shift( self, - periods: int = ..., - freq: Frequency | dt.timedelta | None = ..., + periods: int | Sequence[int] = ..., + freq: DateOffset | dt.timedelta | _str | None = ..., axis: Axis = ..., - fill_value: Hashable | None = ..., + fill_value: Scalar | NAType | None = ..., ) -> Self: ... @overload def set_index( @@ -1989,9 +1989,10 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): self, periods: int = ..., fill_method: None = ..., - limit: int | None = ..., - freq=..., - **kwargs: Any, # TODO: make more precise https://github.com/pandas-dev/pandas-stubs/issues/1169 + freq: DateOffset | dt.timedelta | _str | None = ..., + *, + axis: Axis = ..., + fill_value: Scalar | NAType | None = ..., ) -> Self: ... def pop(self, item: _str) -> Series: ... def pow( diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index c6b7f8d5..ada315b7 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -102,6 +102,7 @@ from pandas._libs.lib import NoDefault from pandas._libs.missing import NAType from pandas._libs.tslibs import BaseOffset from pandas._libs.tslibs.nattype import NaTType +from pandas._libs.tslibs.offsets import DateOffset from pandas._typing import ( S1, S2, @@ -125,7 +126,6 @@ from pandas._typing import ( FilePath, FillnaOptions, FloatDtypeArg, - Frequency, GroupByObjectNonScalar, HashableT1, IgnoreRaise, @@ -1130,10 +1130,10 @@ class Series(IndexOpsMixin[S1], NDFrame): ) -> Series[S1]: ... def shift( self, - periods: int = ..., - freq: Frequency | timedelta | None = ..., - axis: AxisIndex = ..., - fill_value: object | None = ..., + periods: int | Sequence[int] = ..., + freq: DateOffset | timedelta | _str | None = ..., + axis: Axis = ..., + fill_value: Scalar | NAType | None = ..., ) -> UnknownSeries: ... def info( self, @@ -1549,11 +1549,11 @@ class Series(IndexOpsMixin[S1], NDFrame): def pct_change( self, periods: int = ..., - fill_method: _str = ..., - limit: int | None = ..., - freq=..., - **kwargs: Any, # TODO: make more precise https://github.com/pandas-dev/pandas-stubs/issues/1169 - ) -> Series[S1]: ... + fill_method: None = ..., + freq: DateOffset | timedelta | _str | None = ..., + *, + fill_value: Scalar | NAType | None = ..., + ) -> Series[float]: ... def first_valid_index(self) -> Scalar: ... def last_valid_index(self) -> Scalar: ... @overload diff --git a/tests/test_frame.py b/tests/test_frame.py index 99f11b3c..2bde7062 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -2428,7 +2428,13 @@ def test_groupby_series_methods() -> None: def test_dataframe_pct_change() -> None: df = pd.DataFrame({"x": [1, 2, 2, 3, 3], "y": [10, 20, 30, 40, 50]}) - df.pct_change(fill_method=None) + check(assert_type(df.pct_change(), pd.DataFrame), pd.DataFrame) + check(assert_type(df.pct_change(fill_method=None), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.pct_change(axis="columns", periods=-1), pd.DataFrame), + pd.DataFrame, + ) + check(assert_type(df.pct_change(fill_value=0), pd.DataFrame), pd.DataFrame) def test_indexslice_setitem(): diff --git a/tests/test_series.py b/tests/test_series.py index 0a9d6909..cb6fe085 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -441,6 +441,26 @@ def test_types_shift() -> None: check(assert_type(s.shift(freq="1D"), pd.Series), pd.Series, np.integer) +def test_series_pct_change() -> None: + s = pd.Series([1, 2, 3], index=pd.date_range("2020", periods=3)) + check(assert_type(s.pct_change(), "pd.Series[float]"), pd.Series, np.floating) + check( + assert_type(s.pct_change(fill_method=None), "pd.Series[float]"), + pd.Series, + np.floating, + ) + check( + assert_type(s.pct_change(periods=-1), "pd.Series[float]"), + pd.Series, + np.floating, + ) + check( + assert_type(s.pct_change(fill_value=0), "pd.Series[float]"), + pd.Series, + np.floating, + ) + + def test_types_rank() -> None: s = pd.Series([1, 1, 2, 5, 6, np.nan]) s.rank()