diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index e5e91133f..ec25e869b 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -302,4 +302,6 @@ class StyleExportDict(TypedDict, total=False): hide_column_names: bool css: dict[str, str | int] +CalculationMethod = Literal["single", "table"] + __all__ = ["npt", "type_t"] diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 15edfa047..aa04b6342 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -27,6 +27,7 @@ from pandas.core.groupby.generic import ( _DataFrameGroupByScalar, ) from pandas.core.groupby.grouper import Grouper +from pandas.core.indexers import BaseIndexer from pandas.core.indexes.base import Index from pandas.core.indexing import ( _iLocIndexer, @@ -46,6 +47,7 @@ from pandas.core.window.rolling import ( import xarray as xr from pandas._libs.missing import NAType +from pandas._libs.tslibs import BaseOffset from pandas._typing import ( S1, AggFuncTypeBase, @@ -56,6 +58,7 @@ from pandas._typing import ( Axes, Axis, AxisType, + CalculationMethod, ColspaceArgType, CompressionOptions, Dtype, @@ -1429,7 +1432,7 @@ class DataFrame(NDFrame, OpsMixin): self, min_periods: int = ..., axis: AxisType = ..., - method: Literal["single", "table"] = ..., + method: CalculationMethod = ..., ) -> Expanding[DataFrame]: ... @overload def ffill( @@ -1763,7 +1766,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def rolling( self, - window, + window: int | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, @@ -1771,11 +1774,13 @@ class DataFrame(NDFrame, OpsMixin): on: Hashable | None = ..., axis: AxisType = ..., closed: IntervalClosedType | None = ..., + step: int | None = ..., + method: CalculationMethod = ..., ) -> Window[DataFrame]: ... @overload def rolling( self, - window, + window: int | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, @@ -1783,6 +1788,8 @@ class DataFrame(NDFrame, OpsMixin): on: Hashable | None = ..., axis: AxisType = ..., closed: IntervalClosedType | None = ..., + step: int | None = ..., + method: CalculationMethod = ..., ) -> Rolling[DataFrame]: ... def rpow( self, diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 233e11524..6ce270494 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -34,6 +34,7 @@ from pandas.core.groupby.generic import ( _SeriesGroupByNonScalar, _SeriesGroupByScalar, ) +from pandas.core.indexers import BaseIndexer from pandas.core.indexes.accessors import ( CombinedDatetimelikeProperties, PeriodProperties, @@ -62,6 +63,7 @@ from pandas.core.window.rolling import ( import xarray as xr from pandas._libs.missing import NAType +from pandas._libs.tslibs import BaseOffset from pandas._typing import ( S1, AggFuncTypeBase, @@ -71,6 +73,7 @@ from pandas._typing import ( Axes, Axis, AxisType, + CalculationMethod, CompressionOptions, DtypeObj, FilePathOrBuffer, @@ -1318,7 +1321,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): self, min_periods: int = ..., axis: SeriesAxisType = ..., - method: Literal["single", "table"] = ..., + method: CalculationMethod = ..., ) -> Expanding[Series]: ... def floordiv( self, @@ -1496,7 +1499,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def rolling( self, - window, + window: int | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, @@ -1504,11 +1507,13 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): on: _str | None = ..., axis: SeriesAxisType = ..., closed: _str | None = ..., + step: int | None = ..., + method: CalculationMethod = ..., ) -> Window[Series]: ... @overload def rolling( self, - window, + window: int | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, @@ -1516,6 +1521,8 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): on: _str | None = ..., axis: SeriesAxisType = ..., closed: _str | None = ..., + step: int | None = ..., + method: CalculationMethod = ..., ) -> Rolling[Series]: ... def rpow( self, diff --git a/pandas-stubs/core/window/ewm.pyi b/pandas-stubs/core/window/ewm.pyi index 425f2658f..7bb39e6ba 100644 --- a/pandas-stubs/core/window/ewm.pyi +++ b/pandas-stubs/core/window/ewm.pyi @@ -1,7 +1,6 @@ from typing import ( Any, Generic, - Literal, overload, ) @@ -17,6 +16,7 @@ from pandas._typing import ( AggFuncTypeFrame, AggFuncTypeSeriesToFrame, Axis, + CalculationMethod, NDFrameT, TimedeltaConvertibleTypes, WindowingEngine, @@ -36,7 +36,7 @@ class ExponentialMovingWindow(BaseWindow[NDFrameT], Generic[NDFrameT]): ignore_na: bool = ..., axis: Axis = ..., times: str | np.ndarray | Series | None = ..., - method: Literal["single", "table"] = ..., + method: CalculationMethod = ..., ) -> None: ... @overload def aggregate( diff --git a/tests/test_windowing.py b/tests/test_windowing.py index f241bf58a..8f3da3e34 100644 --- a/tests/test_windowing.py +++ b/tests/test_windowing.py @@ -1,4 +1,5 @@ import numpy as np +import pandas as pd from pandas import ( DataFrame, Series, @@ -304,3 +305,35 @@ def test_ewm_aggregate_series() -> None: DataFrame, ) check(assert_type(S.ewm(span=10).agg("sum"), Series), Series) + + +def test_rolling_step_method() -> None: + check( + assert_type(DF.rolling(10, step=5, method="single"), "Rolling[DataFrame]"), + Rolling, + ) + check(assert_type(DF.rolling(10, method="table"), "Rolling[DataFrame]"), Rolling) + + +def test_rolling_window() -> None: + df_time = pd.DataFrame( + {"B": [0, 1, 2, np.nan, 4]}, + index=[ + pd.Timestamp("20130101 09:00:00"), + pd.Timestamp("20130101 09:00:02"), + pd.Timestamp("20130101 09:00:03"), + pd.Timestamp("20130101 09:00:05"), + pd.Timestamp("20130101 09:00:06"), + ], + ) + + indexer = pd.api.indexers.FixedForwardWindowIndexer(window_size=2) + check( + assert_type(df_time.rolling(window=indexer, min_periods=1).sum(), DataFrame), + DataFrame, + ) + s = df_time.iloc[:, 0] + check( + assert_type(s.rolling(window=indexer, min_periods=1).sum(), Series), + Series, + )