ENH: Improve Roling, Expanding and EWM (#297)

* ENH: IMprove Roling, Expanding and EWM

* ENH: Add generic

* TST: Add tests for rolling

* TST: Add tests and fix types

* Improve typing

* TYP Further improvements and refactoring

Co-authored-by: Kevin Sheppard <[email protected]>

Author: bashtage

pandas-stubs/_typing.pyi

@@ -98,12 +98,17 @@ FuncType = Callable[..., Any]
 F = TypeVar("F", bound=FuncType)
 HashableT = TypeVar("HashableT", bound=Hashable)
 
-AggFuncTypeBase = Union[Callable, str]
-AggFuncTypeDict = dict[Hashable, Union[AggFuncTypeBase, list[AggFuncTypeBase]]]
-AggFuncType = Union[
+AggFuncTypeBase = Union[Callable, str, np.ufunc]
+AggFuncTypeDictSeries = dict[Hashable, AggFuncTypeBase]
+AggFuncTypeDictFrame = dict[Hashable, Union[AggFuncTypeBase, list[AggFuncTypeBase]]]
+AggFuncTypeSeriesToFrame = Union[
+    list[AggFuncTypeBase],
+    AggFuncTypeDictSeries,
+]
+AggFuncTypeFrame = Union[
     AggFuncTypeBase,
     list[AggFuncTypeBase],
-    AggFuncTypeDict,
+    AggFuncTypeDictFrame,
 ]
 
 num = complex
@@ -264,6 +269,18 @@ FileWriteMode = Literal[
 ]
 ColspaceArgType = str | int | Sequence[int | str] | Mapping[Hashable, str | int]
 
+# Windowing rank methods
+WindowingRankType = Literal["average", "min", "max"]
+WindowingEngine = Union[Literal["cython", "numba"], None]
+
+class _WindowingNumbaKwargs(TypedDict, total=False):
+    nopython: bool
+    nogil: bool
+    parallel: bool
+
+WindowingEngineKwargs = Union[_WindowingNumbaKwargs, None]
+QuantileInterpolation = Literal["linear", "lower", "higher", "midpoint", "nearest"]
+
 class StyleExportDict(TypedDict, total=False):
     apply: Any
     table_attributes: Any

pandas-stubs/core/frame.pyi

@@ -34,6 +34,10 @@ from pandas.core.indexing import (
 )
 from pandas.core.resample import Resampler
 from pandas.core.series import Series
+from pandas.core.window import (
+    Expanding,
+    ExponentialMovingWindow,
+)
 from pandas.core.window.rolling import (
     Rolling,
     Window,
@@ -43,9 +47,9 @@ import xarray as xr
 from pandas._libs.missing import NAType
 from pandas._typing import (
     S1,
-    AggFuncType,
     AggFuncTypeBase,
-    AggFuncTypeDict,
+    AggFuncTypeDictFrame,
+    AggFuncTypeFrame,
     AnyArrayLike,
     ArrayLike,
     Axes,
@@ -65,6 +69,7 @@ from pandas._typing import (
     IndexingInt,
     IndexLabel,
     IndexType,
+    IntervalClosedType,
     JsonFrameOrient,
     Label,
     Level,
@@ -74,6 +79,7 @@ from pandas._typing import (
     MergeHow,
     NaPosition,
     ParquetEngine,
+    QuantileInterpolation,
     ReadBuffer,
     Renamer,
     ReplaceMethod,
@@ -1052,7 +1058,7 @@ class DataFrame(NDFrame, OpsMixin):
     @overload
     def agg(
         self,
-        func: list[AggFuncTypeBase] | AggFuncTypeDict = ...,
+        func: list[AggFuncTypeBase] | AggFuncTypeDictFrame = ...,
         axis: AxisType = ...,
         **kwargs,
     ) -> DataFrame: ...
@@ -1063,13 +1069,13 @@ class DataFrame(NDFrame, OpsMixin):
     @overload
     def aggregate(
         self,
-        func: list[AggFuncTypeBase] | AggFuncTypeDict,
+        func: list[AggFuncTypeBase] | AggFuncTypeDictFrame,
         axis: AxisType = ...,
         **kwargs,
     ) -> DataFrame: ...
     def transform(
         self,
-        func: AggFuncType,
+        func: AggFuncTypeFrame,
         axis: AxisType = ...,
         *args,
         **kwargs,
@@ -1165,17 +1171,15 @@ class DataFrame(NDFrame, OpsMixin):
         q: float = ...,
         axis: AxisType = ...,
         numeric_only: _bool = ...,
-        interpolation: _str
-        | Literal["linear", "lower", "higher", "midpoint", "nearest"] = ...,
+        interpolation: QuantileInterpolation = ...,
     ) -> Series: ...
     @overload
     def quantile(
         self,
         q: list[float] | np.ndarray,
         axis: AxisType = ...,
         numeric_only: _bool = ...,
-        interpolation: _str
-        | Literal["linear", "lower", "higher", "midpoint", "nearest"] = ...,
+        interpolation: QuantileInterpolation = ...,
     ) -> DataFrame: ...
     def to_timestamp(
         self,
@@ -1412,8 +1416,13 @@ class DataFrame(NDFrame, OpsMixin):
         adjust: _bool = ...,
         ignore_na: _bool = ...,
         axis: AxisType = ...,
-    ) -> DataFrame: ...
-    def expanding(self, min_periods: int = ..., axis: AxisType = ...): ...  # for now
+    ) -> ExponentialMovingWindow[DataFrame]: ...
+    def expanding(
+        self,
+        min_periods: int = ...,
+        axis: AxisType = ...,
+        method: Literal["single", "table"] = ...,
+    ) -> Expanding[DataFrame]: ...
     @overload
     def ffill(
         self,
@@ -1767,21 +1776,21 @@ class DataFrame(NDFrame, OpsMixin):
         center: _bool = ...,
         *,
         win_type: _str,
-        on: _str | None = ...,
+        on: Hashable | None = ...,
         axis: AxisType = ...,
-        closed: _str | None = ...,
-    ) -> Window: ...
+        closed: IntervalClosedType | None = ...,
+    ) -> Window[DataFrame]: ...
     @overload
     def rolling(
         self,
         window,
         min_periods: int | None = ...,
         center: _bool = ...,
         *,
-        on: _str | None = ...,
+        on: Hashable | None = ...,
         axis: AxisType = ...,
-        closed: _str | None = ...,
-    ) -> Rolling: ...
+        closed: IntervalClosedType | None = ...,
+    ) -> Rolling[DataFrame]: ...
     def rpow(
         self,
         other,

pandas-stubs/core/groupby/generic.pyi

@@ -26,8 +26,8 @@ from pandas.core.series import Series
 
 from pandas._typing import (
     S1,
-    AggFuncType,
     AggFuncTypeBase,
+    AggFuncTypeFrame,
     AxisType,
     Level,
     ListLike,
@@ -154,8 +154,8 @@ class DataFrameGroupBy(GroupBy):
     def apply(  # pyright: ignore[reportOverlappingOverload]
         self, func: Callable[[Iterable], float], *args, **kwargs
     ) -> DataFrame: ...
-    def aggregate(self, arg: AggFuncType = ..., *args, **kwargs) -> DataFrame: ...
-    def agg(self, arg: AggFuncType = ..., *args, **kwargs) -> DataFrame: ...
+    def aggregate(self, arg: AggFuncTypeFrame = ..., *args, **kwargs) -> DataFrame: ...
+    agg = aggregate
     def transform(self, func, *args, **kwargs): ...
     def filter(
         self, func: Callable, dropna: bool = ..., *args, **kwargs

pandas-stubs/core/series.pyi

@@ -49,7 +49,11 @@ from pandas.core.indexing import (
 )
 from pandas.core.resample import Resampler
 from pandas.core.strings import StringMethods
-from pandas.core.window import ExponentialMovingWindow
+from pandas.core.window import (
+    Expanding,
+    ExponentialMovingWindow,
+    Rolling,
+)
 from pandas.core.window.rolling import (
     Rolling,
     Window,
@@ -60,7 +64,8 @@ from pandas._libs.missing import NAType
 from pandas._typing import (
     S1,
     AggFuncTypeBase,
-    AggFuncTypeDict,
+    AggFuncTypeDictFrame,
+    AggFuncTypeSeriesToFrame,
     ArrayLike,
     Axes,
     Axis,
@@ -80,6 +85,7 @@ from pandas._typing import (
     MaskType,
     MergeHow,
     NaPosition,
+    QuantileInterpolation,
     Renamer,
     ReplaceMethod,
     Scalar,
@@ -454,15 +460,13 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     def quantile(
         self,
         q: float = ...,
-        interpolation: _str
-        | Literal["linear", "lower", "higher", "midpoint", "nearest"] = ...,
+        interpolation: QuantileInterpolation = ...,
     ) -> float: ...
     @overload
     def quantile(
         self,
         q: _ListLike,
-        interpolation: _str
-        | Literal["linear", "lower", "higher", "midpoint", "nearest"] = ...,
+        interpolation: QuantileInterpolation = ...,
     ) -> Series[S1]: ...
     def corr(
         self,
@@ -629,27 +633,12 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     @overload
     def aggregate(
         self,
-        func: list[AggFuncTypeBase] | dict[Hashable, AggFuncTypeBase] = ...,
-        axis: SeriesAxisType = ...,
-        *args,
-        **kwargs,
-    ) -> Series[S1]: ...
-    @overload
-    def agg(
-        self,
-        func: AggFuncTypeBase,
-        axis: SeriesAxisType = ...,
-        *args,
-        **kwargs,
-    ) -> S1: ...
-    @overload
-    def agg(
-        self,
-        func: list[AggFuncTypeBase] | dict[Hashable, AggFuncTypeBase] = ...,
+        func: AggFuncTypeSeriesToFrame = ...,
         axis: SeriesAxisType = ...,
         *args,
         **kwargs,
     ) -> Series[S1]: ...
+    agg = aggregate
     @overload
     def transform(
         self,
@@ -661,7 +650,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     @overload
     def transform(
         self,
-        func: list[AggFuncTypeBase] | AggFuncTypeDict,
+        func: list[AggFuncTypeBase] | AggFuncTypeDictFrame,
         axis: SeriesAxisType = ...,
         *args,
         **kwargs,
@@ -1322,10 +1311,13 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         adjust: _bool = ...,
         ignore_na: _bool = ...,
         axis: SeriesAxisType = ...,
-    ) -> ExponentialMovingWindow: ...
+    ) -> ExponentialMovingWindow[Series]: ...
     def expanding(
-        self, min_periods: int = ..., axis: SeriesAxisType = ...
-    ) -> DataFrame: ...
+        self,
+        min_periods: int = ...,
+        axis: SeriesAxisType = ...,
+        method: Literal["single", "table"] = ...,
+    ) -> Expanding[Series]: ...
     def floordiv(
         self,
         other: num | _ListLike | Series[S1],
@@ -1528,7 +1520,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         on: _str | None = ...,
         axis: SeriesAxisType = ...,
         closed: _str | None = ...,
-    ) -> Window: ...
+    ) -> Window[Series]: ...
     @overload
     def rolling(
         self,
@@ -1539,7 +1531,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         on: _str | None = ...,
         axis: SeriesAxisType = ...,
         closed: _str | None = ...,
-    ) -> Rolling: ...
+    ) -> Rolling[Series]: ...
     def rpow(
         self,
         other: Series[S1] | Scalar,