Allow Timedelta for Series.rolling and DataFrame.rolling (#705)

* Add Timedelta annotation to .rolling

* Allow datetime.timedelta as well

* Reorder and indent tests

Author: pinealan

pandas-stubs/core/frame.pyi

@@ -1955,7 +1955,7 @@ class DataFrame(NDFrame, OpsMixin):
     @overload
     def rolling(
         self,
-        window: int | str | BaseOffset | BaseIndexer,
+        window: int | str | _dt.timedelta | BaseOffset | BaseIndexer,
         min_periods: int | None = ...,
         center: _bool = ...,
         on: Hashable | None = ...,
@@ -1969,7 +1969,7 @@ class DataFrame(NDFrame, OpsMixin):
     @overload
     def rolling(
         self,
-        window: int | str | BaseOffset | BaseIndexer,
+        window: int | str | _dt.timedelta | BaseOffset | BaseIndexer,
         min_periods: int | None = ...,
         center: _bool = ...,
         on: Hashable | None = ...,

pandas-stubs/core/series.pyi

@@ -1834,7 +1834,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     @overload
     def rolling(
         self,
-        window: int | _str | BaseOffset | BaseIndexer,
+        window: int | _str | timedelta | BaseOffset | BaseIndexer,
         min_periods: int | None = ...,
         center: _bool = ...,
         on: _str | None = ...,
@@ -1848,7 +1848,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     @overload
     def rolling(
         self,
-        window: int | _str | BaseOffset | BaseIndexer,
+        window: int | _str | timedelta | BaseOffset | BaseIndexer,
         min_periods: int | None = ...,
         center: _bool = ...,
         on: _str | None = ...,

tests/test_windowing.py

@@ -1,8 +1,11 @@
+import datetime as dt
+
 import numpy as np
 import pandas as pd
 from pandas import (
     DataFrame,
     Series,
+    Timedelta,
     date_range,
 )
 from pandas.core.window import (
@@ -62,6 +65,14 @@ def test_rolling_datetime_index() -> None:
         Series,
     )
 
+    td = Timedelta("1D")
+    check(assert_type(DF_DTI.rolling(td), "Rolling[DataFrame]"), Rolling, DataFrame)
+    check(assert_type(S_DTI.rolling(td), "Rolling[Series]"), Rolling, Series)
+
+    dttd = dt.timedelta(days=1)
+    check(assert_type(DF_DTI.rolling(dttd), "Rolling[DataFrame]"), Rolling, DataFrame)
+    check(assert_type(S_DTI.rolling(dttd), "Rolling[Series]"), Rolling, Series)
+
 
 def test_rolling_apply() -> None:
     check(assert_type(DF.rolling(10).apply(np.mean), DataFrame), DataFrame)