diff --git a/.gitignore b/.gitignore index 88ebdb2d7..f47c62c6e 100644 --- a/.gitignore +++ b/.gitignore @@ -120,6 +120,9 @@ venv.bak/ # PyCharm project settings .idea/ +# VSCode project settings +.vscode/ + # mkdocs documentation /site diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index cf825a2f3..fcd6b996a 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1776,7 +1776,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def rolling( self, - window: int | BaseOffset | BaseIndexer, + window: int | str | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, @@ -1790,7 +1790,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def rolling( self, - window: int | BaseOffset | BaseIndexer, + window: int | str | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index 94161ed2d..0a12a9a1b 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -86,6 +86,7 @@ from pandas._typing import ( HashableT3, IgnoreRaise, IndexingInt, + IntervalClosedType, JoinHow, JsonSeriesOrient, Level, @@ -1574,28 +1575,28 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def rolling( self, - window: int | BaseOffset | BaseIndexer, + window: int | _str | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, win_type: _str, on: _str | None = ..., axis: SeriesAxisType = ..., - closed: _str | None = ..., + closed: IntervalClosedType | None = ..., step: int | None = ..., method: CalculationMethod = ..., ) -> Window[Series]: ... @overload def rolling( self, - window: int | BaseOffset | BaseIndexer, + window: int | _str | BaseOffset | BaseIndexer, min_periods: int | None = ..., center: _bool = ..., *, win_type: None = ..., on: _str | None = ..., axis: SeriesAxisType = ..., - closed: _str | None = ..., + closed: IntervalClosedType | None = ..., step: int | None = ..., method: CalculationMethod = ..., ) -> Rolling[Series]: ... diff --git a/tests/test_windowing.py b/tests/test_windowing.py index 8f3da3e34..bf0d88dbf 100644 --- a/tests/test_windowing.py +++ b/tests/test_windowing.py @@ -13,9 +13,13 @@ from tests import check +from pandas.tseries.frequencies import to_offset + IDX = date_range("1/1/2000", periods=700, freq="D") S = Series(np.random.standard_normal(700)) DF = DataFrame({"col1": S, "col2": S}) +S_DTI = Series(data=np.random.standard_normal(700), index=IDX) +DF_DTI = DataFrame(data=np.random.standard_normal(700), index=IDX) def test_rolling_basic() -> None: @@ -43,6 +47,22 @@ def test_rolling_basic_math() -> None: check(assert_type(DF.rolling(10).rank("max"), DataFrame), DataFrame) +def test_rolling_datetime_index() -> None: + offset_1d = to_offset("1D") + assert offset_1d is not None + + check(assert_type(DF_DTI.rolling("1D"), "Rolling[DataFrame]"), Rolling, DataFrame) + check( + assert_type(DF_DTI.rolling(offset_1d), "Rolling[DataFrame]"), Rolling, DataFrame + ) + check(assert_type(S_DTI.rolling("1D"), "Rolling[Series]"), Rolling, Series) + check( + assert_type(S_DTI.rolling(offset_1d), "Rolling[Series]"), + Rolling, + Series, + ) + + def test_rolling_apply() -> None: check(assert_type(DF.rolling(10).apply(np.mean), DataFrame), DataFrame)