GH1169 Improve parameter types for DataFrame.pct_change and Series.pct_change (#1194)

* type pct_change kwargs according to shift params

* remove duplicate params and add defaults

* remove typeddict

* address comments

* add tests, fix axis argument type

* fix shift/pct_change param types

* fix return type of series pct_change, add test for series.pct_change

Author: yangdanny97

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(

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

tests/test_frame.py

@@ -2437,7 +2437,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():

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()