From 1b574c1d158a6b00268183328cbd3e6290dc8132 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Mon, 26 Sep 2022 22:15:40 +0100 Subject: [PATCH 1/3] ENH: Add step to rolling --- pandas-stubs/_typing.pyi | 2 ++ pandas-stubs/core/frame.pyi | 11 ++++++++--- pandas-stubs/core/series.pyi | 11 ++++++++--- pandas-stubs/core/window/ewm.pyi | 4 ++-- tests/test_windowing.py | 8 ++++++++ 5 files changed, 28 insertions(+), 8 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index e5e91133f..596624184 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] +RollingMethod = Literal["single", "table"] + __all__ = ["npt", "type_t"] diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 15edfa047..4776216e7 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -84,6 +84,7 @@ from pandas._typing import ( ReadBuffer, Renamer, ReplaceMethod, + RollingMethod, Scalar, ScalarT, SeriesAxisType, @@ -1429,7 +1430,7 @@ class DataFrame(NDFrame, OpsMixin): self, min_periods: int = ..., axis: AxisType = ..., - method: Literal["single", "table"] = ..., + method: RollingMethod = ..., ) -> Expanding[DataFrame]: ... @overload def ffill( @@ -1763,7 +1764,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def rolling( self, - window, + window: int, min_periods: int | None = ..., center: _bool = ..., *, @@ -1771,11 +1772,13 @@ class DataFrame(NDFrame, OpsMixin): on: Hashable | None = ..., axis: AxisType = ..., closed: IntervalClosedType | None = ..., + step: int | None = ..., + method: RollingMethod = ..., ) -> Window[DataFrame]: ... @overload def rolling( self, - window, + window: int, min_periods: int | None = ..., center: _bool = ..., *, @@ -1783,6 +1786,8 @@ class DataFrame(NDFrame, OpsMixin): on: Hashable | None = ..., axis: AxisType = ..., closed: IntervalClosedType | None = ..., + step: int | None = ..., + method: RollingMethod = ..., ) -> Rolling[DataFrame]: ... def rpow( self, diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 233e11524..5f3a4a136 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -88,6 +88,7 @@ from pandas._typing import ( QuantileInterpolation, Renamer, ReplaceMethod, + RollingMethod, Scalar, SeriesAxisType, SortKind, @@ -1318,7 +1319,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): self, min_periods: int = ..., axis: SeriesAxisType = ..., - method: Literal["single", "table"] = ..., + method: RollingMethod = ..., ) -> Expanding[Series]: ... def floordiv( self, @@ -1496,7 +1497,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def rolling( self, - window, + window: int, min_periods: int | None = ..., center: _bool = ..., *, @@ -1504,11 +1505,13 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): on: _str | None = ..., axis: SeriesAxisType = ..., closed: _str | None = ..., + step: int | None = ..., + method: RollingMethod = ..., ) -> Window[Series]: ... @overload def rolling( self, - window, + window: int, min_periods: int | None = ..., center: _bool = ..., *, @@ -1516,6 +1519,8 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): on: _str | None = ..., axis: SeriesAxisType = ..., closed: _str | None = ..., + step: int | None = ..., + method: RollingMethod = ..., ) -> Rolling[Series]: ... def rpow( self, diff --git a/pandas-stubs/core/window/ewm.pyi b/pandas-stubs/core/window/ewm.pyi index 425f2658f..025b425b7 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, ) @@ -18,6 +17,7 @@ from pandas._typing import ( AggFuncTypeSeriesToFrame, Axis, NDFrameT, + RollingMethod, TimedeltaConvertibleTypes, WindowingEngine, WindowingEngineKwargs, @@ -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: RollingMethod = ..., ) -> None: ... @overload def aggregate( diff --git a/tests/test_windowing.py b/tests/test_windowing.py index f241bf58a..556866b4b 100644 --- a/tests/test_windowing.py +++ b/tests/test_windowing.py @@ -304,3 +304,11 @@ def test_ewm_aggregate_series() -> None: DataFrame, ) check(assert_type(S.ewm(span=10).agg("sum"), Series), Series) + + +def test_rolling_step_method(): + check( + assert_type(DF.rolling(10, step=5, method="single"), "Rolling[DataFrame]"), + Rolling, + ) + check(assert_type(DF.rolling(10, method="table"), "Rolling[DataFrame]"), Rolling) From ac910c5b69a01a7139c0eeb78a6bb8aeb1e93f42 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Tue, 27 Sep 2022 09:05:58 +0100 Subject: [PATCH 2/3] CLN: Rename RollingMethod to CalculationMethod --- pandas-stubs/_typing.pyi | 2 +- pandas-stubs/core/frame.pyi | 8 ++++---- pandas-stubs/core/series.pyi | 8 ++++---- pandas-stubs/core/window/ewm.pyi | 4 ++-- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index 596624184..ec25e869b 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -302,6 +302,6 @@ class StyleExportDict(TypedDict, total=False): hide_column_names: bool css: dict[str, str | int] -RollingMethod = Literal["single", "table"] +CalculationMethod = Literal["single", "table"] __all__ = ["npt", "type_t"] diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 4776216e7..34515eaf8 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -56,6 +56,7 @@ from pandas._typing import ( Axes, Axis, AxisType, + CalculationMethod, ColspaceArgType, CompressionOptions, Dtype, @@ -84,7 +85,6 @@ from pandas._typing import ( ReadBuffer, Renamer, ReplaceMethod, - RollingMethod, Scalar, ScalarT, SeriesAxisType, @@ -1430,7 +1430,7 @@ class DataFrame(NDFrame, OpsMixin): self, min_periods: int = ..., axis: AxisType = ..., - method: RollingMethod = ..., + method: CalculationMethod = ..., ) -> Expanding[DataFrame]: ... @overload def ffill( @@ -1773,7 +1773,7 @@ class DataFrame(NDFrame, OpsMixin): axis: AxisType = ..., closed: IntervalClosedType | None = ..., step: int | None = ..., - method: RollingMethod = ..., + method: CalculationMethod = ..., ) -> Window[DataFrame]: ... @overload def rolling( @@ -1787,7 +1787,7 @@ class DataFrame(NDFrame, OpsMixin): axis: AxisType = ..., closed: IntervalClosedType | None = ..., step: int | None = ..., - method: RollingMethod = ..., + method: CalculationMethod = ..., ) -> Rolling[DataFrame]: ... def rpow( self, diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 5f3a4a136..9ade2f510 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -71,6 +71,7 @@ from pandas._typing import ( Axes, Axis, AxisType, + CalculationMethod, CompressionOptions, DtypeObj, FilePathOrBuffer, @@ -88,7 +89,6 @@ from pandas._typing import ( QuantileInterpolation, Renamer, ReplaceMethod, - RollingMethod, Scalar, SeriesAxisType, SortKind, @@ -1319,7 +1319,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): self, min_periods: int = ..., axis: SeriesAxisType = ..., - method: RollingMethod = ..., + method: CalculationMethod = ..., ) -> Expanding[Series]: ... def floordiv( self, @@ -1506,7 +1506,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): axis: SeriesAxisType = ..., closed: _str | None = ..., step: int | None = ..., - method: RollingMethod = ..., + method: CalculationMethod = ..., ) -> Window[Series]: ... @overload def rolling( @@ -1520,7 +1520,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): axis: SeriesAxisType = ..., closed: _str | None = ..., step: int | None = ..., - method: RollingMethod = ..., + method: CalculationMethod = ..., ) -> Rolling[Series]: ... def rpow( self, diff --git a/pandas-stubs/core/window/ewm.pyi b/pandas-stubs/core/window/ewm.pyi index 025b425b7..7bb39e6ba 100644 --- a/pandas-stubs/core/window/ewm.pyi +++ b/pandas-stubs/core/window/ewm.pyi @@ -16,8 +16,8 @@ from pandas._typing import ( AggFuncTypeFrame, AggFuncTypeSeriesToFrame, Axis, + CalculationMethod, NDFrameT, - RollingMethod, TimedeltaConvertibleTypes, WindowingEngine, WindowingEngineKwargs, @@ -36,7 +36,7 @@ class ExponentialMovingWindow(BaseWindow[NDFrameT], Generic[NDFrameT]): ignore_na: bool = ..., axis: Axis = ..., times: str | np.ndarray | Series | None = ..., - method: RollingMethod = ..., + method: CalculationMethod = ..., ) -> None: ... @overload def aggregate( From d1e3581474aed9df4d3246b391778dbbfe3023b0 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Tue, 27 Sep 2022 17:09:13 +0100 Subject: [PATCH 3/3] ENH/TST: Add other window types and a test --- pandas-stubs/core/frame.pyi | 6 ++++-- pandas-stubs/core/series.pyi | 6 ++++-- tests/test_windowing.py | 27 ++++++++++++++++++++++++++- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 34515eaf8..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, @@ -1764,7 +1766,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def rolling( self, - window: int, + window: int | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, @@ -1778,7 +1780,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def rolling( self, - window: int, + window: int | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 9ade2f510..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, @@ -1497,7 +1499,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def rolling( self, - window: int, + window: int | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, @@ -1511,7 +1513,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def rolling( self, - window: int, + window: int | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, diff --git a/tests/test_windowing.py b/tests/test_windowing.py index 556866b4b..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, @@ -306,9 +307,33 @@ def test_ewm_aggregate_series() -> None: check(assert_type(S.ewm(span=10).agg("sum"), Series), Series) -def test_rolling_step_method(): +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, + )