diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index f67b3f908..888141df1 100644 --- a/pandas-stubs/_typing.pyi +++ b/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 diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index ee625bbef..4569d3ba7 100644 --- a/pandas-stubs/core/frame.pyi +++ b/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, @@ -42,9 +46,9 @@ import xarray as xr from pandas._typing import ( S1, - AggFuncType, AggFuncTypeBase, - AggFuncTypeDict, + AggFuncTypeDictFrame, + AggFuncTypeFrame, AnyArrayLike, ArrayLike, Axes, @@ -64,6 +68,7 @@ from pandas._typing import ( IndexingInt, IndexLabel, IndexType, + IntervalClosedType, JsonFrameOrient, Label, Level, @@ -73,6 +78,7 @@ from pandas._typing import ( MergeHow, NaPosition, ParquetEngine, + QuantileInterpolation, ReadBuffer, Renamer, ReplaceMethod, @@ -1050,7 +1056,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def agg( self, - func: list[AggFuncTypeBase] | AggFuncTypeDict = ..., + func: list[AggFuncTypeBase] | AggFuncTypeDictFrame = ..., axis: AxisType = ..., **kwargs, ) -> DataFrame: ... @@ -1061,13 +1067,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, @@ -1163,8 +1169,7 @@ 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( @@ -1172,8 +1177,7 @@ class DataFrame(NDFrame, OpsMixin): 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, @@ -1410,8 +1414,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, @@ -1765,10 +1774,10 @@ 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, @@ -1776,10 +1785,10 @@ class DataFrame(NDFrame, OpsMixin): 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, diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index 8648d0fc1..e11354529 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/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 diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index e64e56fc4..62d91f6c5 100644 --- a/pandas-stubs/core/series.pyi +++ b/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, @@ -59,7 +63,8 @@ import xarray as xr from pandas._typing import ( S1, AggFuncTypeBase, - AggFuncTypeDict, + AggFuncTypeDictFrame, + AggFuncTypeSeriesToFrame, ArrayLike, Axes, Axis, @@ -79,6 +84,7 @@ from pandas._typing import ( MaskType, MergeHow, NaPosition, + QuantileInterpolation, Renamer, ReplaceMethod, Scalar, @@ -453,15 +459,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, @@ -628,27 +632,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, @@ -660,7 +649,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def transform( self, - func: list[AggFuncTypeBase] | AggFuncTypeDict, + func: list[AggFuncTypeBase] | AggFuncTypeDictFrame, axis: SeriesAxisType = ..., *args, **kwargs, @@ -1321,10 +1310,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], @@ -1527,7 +1519,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): on: _str | None = ..., axis: SeriesAxisType = ..., closed: _str | None = ..., - ) -> Window: ... + ) -> Window[Series]: ... @overload def rolling( self, @@ -1538,7 +1530,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, diff --git a/pandas-stubs/core/window/ewm.pyi b/pandas-stubs/core/window/ewm.pyi index 42afcebc2..ce70571d0 100644 --- a/pandas-stubs/core/window/ewm.pyi +++ b/pandas-stubs/core/window/ewm.pyi @@ -1,28 +1,86 @@ -from pandas.core.window.rolling import _Rolling +from typing import ( + Any, + Generic, + Literal, + overload, +) -class ExponentialMovingWindow(_Rolling): - obj = ... - com = ... - min_periods: int = ... - adjust = ... - ignore_na = ... - axis = ... - on = ... +import numpy as np +from pandas import ( + DataFrame, + Series, +) +from pandas.core.window.rolling import BaseWindow + +from pandas._typing import ( + AggFuncTypeBase, + AggFuncTypeFrame, + AggFuncTypeSeriesToFrame, + Axis, + NDFrameT, + TimedeltaConvertibleTypes, + WindowingEngine, + WindowingEngineKwargs, +) + +class ExponentialMovingWindow(BaseWindow[NDFrameT], Generic[NDFrameT]): def __init__( self, - obj, - com=..., - span=..., - halflife=..., - alpha=..., - min_periods: int = ..., + obj: NDFrameT, + com: float | None = ..., + span: float | None = ..., + halflife: TimedeltaConvertibleTypes | None = ..., + alpha: float | None = ..., + min_periods: int | None = ..., adjust: bool = ..., ignore_na: bool = ..., - axis: int = ..., + axis: Axis = ..., + times: str | np.ndarray | Series | None = ..., + method: Literal["single", "table"] = ..., ) -> None: ... - def mean(self, *args, **kwargs): ... - def std(self, bias: bool = ..., *args, **kwargs): ... - vol = ... - def var(self, bias: bool = ..., *args, **kwargs): ... - def cov(self, other=..., pairwise=..., bias: bool = ..., **kwargs): ... - def corr(self, other=..., pairwise=..., **kwargs): ... + @overload + def aggregate( + self: ExponentialMovingWindow[Series], + func: AggFuncTypeBase, + *args: Any, + **kwargs: Any, + ) -> Series: ... + @overload + def aggregate( + self: ExponentialMovingWindow[Series], + func: AggFuncTypeSeriesToFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + @overload + def aggregate( + self: ExponentialMovingWindow[DataFrame], + func: AggFuncTypeFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + def mean( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def sum( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def std(self, bias: bool = ...) -> NDFrameT: ... + def var(self, bias: bool = ...) -> NDFrameT: ... + def cov( + self, + other: DataFrame | Series | None = ..., + pairwise: bool | None = ..., + bias: bool = ..., + ) -> NDFrameT: ... + def corr( + self, + other: DataFrame | Series | None = ..., + pairwise: bool | None = ..., + ) -> NDFrameT: ... diff --git a/pandas-stubs/core/window/expanding.pyi b/pandas-stubs/core/window/expanding.pyi index cf67b0bef..9139c74bd 100644 --- a/pandas-stubs/core/window/expanding.pyi +++ b/pandas-stubs/core/window/expanding.pyi @@ -1,27 +1,127 @@ from typing import ( Any, Callable, + overload, ) +import numpy as np from pandas import ( DataFrame, Series, ) -from pandas.core.window.rolling import _Rolling_and_Expanding +from pandas.core.window.rolling import ( + BaseWindowGroupby, + RollingAndExpandingMixin, +) + +from pandas._typing import ( + AggFuncTypeBase, + AggFuncTypeFrame, + AggFuncTypeSeriesToFrame, + NDFrameT, + QuantileInterpolation, + WindowingEngine, + WindowingEngineKwargs, + WindowingRankType, +) -class Expanding(_Rolling_and_Expanding): - def __init__( - self, obj, min_periods: int = ..., center: bool = ..., axis: int = ..., **kwargs - ) -> None: ... - def count(self, **kwargs) -> DataFrame | Series: ... +class Expanding(RollingAndExpandingMixin[NDFrameT]): + @overload + def aggregate( + self: Expanding[Series], func: AggFuncTypeBase, *args: Any, **kwargs: Any + ) -> Series: ... + @overload + def aggregate( + self: Expanding[Series], + func: AggFuncTypeSeriesToFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + @overload + def aggregate( + self: Expanding[DataFrame], + func: AggFuncTypeFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + def count(self) -> NDFrameT: ... def apply( self, func: Callable[..., Any], raw: bool = ..., - engine: str | None = ..., - engine_kwargs: dict[str, bool] | None = ..., + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., args: tuple[Any, ...] | None = ..., kwargs: dict[str, Any] | None = ..., - ): ... + ) -> NDFrameT: ... + def sum( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def max( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def min( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def mean( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def median( + self, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def std( + self, + ddof: int = ..., + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def var( + self, + ddof: int = ..., + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def sem(self, ddof: int = ...) -> NDFrameT: ... + def skew(self) -> NDFrameT: ... + def kurt(self) -> NDFrameT: ... + def quantile( + self, + quantile: float, + interpolation: QuantileInterpolation = ..., + ) -> NDFrameT: ... + def rank( + self, + method: WindowingRankType = ..., + ascending: bool = ..., + pct: bool = ..., + ) -> NDFrameT: ... + def cov( + self, + other: DataFrame | Series | None = ..., + pairwise: bool | None = ..., + ddof: int = ..., + ) -> NDFrameT: ... + def corr( + self, + other: DataFrame | Series | None = ..., + pairwise: bool | None = ..., + ddof: int = ..., + ) -> NDFrameT: ... -class ExpandingGroupby(Expanding): ... +class ExpandingGroupby(BaseWindowGroupby, Expanding): ... diff --git a/pandas-stubs/core/window/rolling.pyi b/pandas-stubs/core/window/rolling.pyi index 9ab886d99..c199bffb8 100644 --- a/pandas-stubs/core/window/rolling.pyi +++ b/pandas-stubs/core/window/rolling.pyi @@ -1,116 +1,185 @@ +from typing import ( + Any, + Callable, + Generic, + overload, +) + import numpy as np from pandas import ( DataFrame, - Index, Series, ) -from pandas.core.base import ( - PandasObject, - SelectionMixin, -) +from pandas.core.base import SelectionMixin from pandas._typing import ( - AggFuncType, - Axis, - Scalar, + AggFuncTypeBase, + AggFuncTypeFrame, + AggFuncTypeSeriesToFrame, + NDFrameT, + QuantileInterpolation, + WindowingEngine, + WindowingEngineKwargs, + WindowingRankType, ) -class _Window(PandasObject, SelectionMixin): - exclusions: set[str] = ... - obj = ... - on = ... - closed = ... - window = ... - min_periods: int = ... - center = ... - win_type: str = ... - win_freq = ... - axis = ... - def __init__( - self, - obj, - window=..., - min_periods: int | None = ..., - center: bool | None = ..., - win_type: str | None = ..., - axis: Axis = ..., - on: str | Index | None = ..., - closed: str | None = ..., - **kwargs, - ) -> None: ... - @property - def is_datetimelike(self) -> bool | None: ... - @property - def is_freq_type(self) -> bool: ... +class BaseWindow(SelectionMixin[NDFrameT], Generic[NDFrameT]): def __getattr__(self, attr: str): ... def __iter__(self): ... + @overload + def aggregate( + self: BaseWindow[Series], func: AggFuncTypeBase, *args: Any, **kwargs: Any + ) -> Series: ... + @overload + def aggregate( + self: BaseWindow[Series], + func: AggFuncTypeSeriesToFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + @overload def aggregate( - self, func: AggFuncType = ..., *args, **kwargs - ) -> Scalar | DataFrame | Series: ... - def agg( - self, func: AggFuncType = ..., *args, **kwargs - ) -> Scalar | DataFrame | Series: ... + self: BaseWindow[DataFrame], + func: AggFuncTypeFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + agg = aggregate -class Window(_Window): - def sum(self, *args, **kwargs): ... - def mean(self, *args, **kwargs): ... - def var(self, ddof: int = ..., *args, **kwargs): ... - def std(self, ddof: int = ..., *args, **kwargs): ... +class BaseWindowGroupby(BaseWindow[NDFrameT]): ... -class _Rolling(_Window): ... +class Window(BaseWindow[NDFrameT]): + @overload + def aggregate( + self: Window[Series], func: AggFuncTypeBase, *args: Any, **kwargs: Any + ) -> Series: ... + @overload + def aggregate( + self: Window[Series], + func: AggFuncTypeSeriesToFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + @overload + def aggregate( + self: Window[DataFrame], + func: AggFuncTypeFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + def sum(self, **kwargs: Any) -> NDFrameT: ... + def mean(self, **kwargs: Any) -> NDFrameT: ... + def var(self, ddof: int = ..., **kwargs: Any) -> NDFrameT: ... + def std(self, ddof: int = ..., **kwargs: Any) -> NDFrameT: ... -class _Rolling_and_Expanding(_Rolling): - def count(self) -> DataFrame | Series: ... +class RollingAndExpandingMixin(BaseWindow[NDFrameT], Generic[NDFrameT]): + def count(self) -> NDFrameT: ... def apply( self, - func, + func: Callable[..., Any], raw: bool = ..., - engine: str = ..., - engine_kwargs: dict | None = ..., - args: tuple | None = ..., - kwargs: dict | None = ..., - ): ... - def sum(self, *args, **kwargs) -> DataFrame | Series: ... - def max(self, *args, **kwargs) -> DataFrame | Series: ... - def min(self, *args, **kwargs) -> DataFrame | Series: ... - def mean(self, *args, **kwargs) -> DataFrame | Series: ... - def median(self, **kwargs) -> DataFrame | Series: ... - def std(self, ddof: int = ..., *args, **kwargs) -> DataFrame | Series: ... - def var(self, ddof: int = ..., *args, **kwargs) -> DataFrame | Series: ... - def skew(self, **kwargs) -> DataFrame | Series: ... - def kurt(self, **kwargs) -> DataFrame | Series: ... + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + args: tuple[Any, ...] | None = ..., + kwargs: dict[str, Any] | None = ..., + ) -> NDFrameT: ... + def sum( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def max( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def min( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def mean( + self, + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def median( + self, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def std( + self, + ddof: int = ..., + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def var( + self, + ddof: int = ..., + *, + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs = ..., + ) -> NDFrameT: ... + def skew(self) -> NDFrameT: ... + def sem(self, ddof: int = ...) -> NDFrameT: ... + def kurt(self) -> NDFrameT: ... def quantile( - self, quantile: float, interpolation: str = ..., **kwargs - ) -> DataFrame | Series: ... + self, + quantile: float, + interpolation: QuantileInterpolation = ..., + ) -> NDFrameT: ... + def rank( + self, + method: WindowingRankType = ..., + ascending: bool = ..., + pct: bool = ..., + ) -> NDFrameT: ... def cov( self, - other: DataFrame | Series | np.ndarray | None = ..., + other: DataFrame | Series | None = ..., pairwise: bool | None = ..., ddof: int = ..., - **kwargs, - ) -> DataFrame | Series: ... + ) -> NDFrameT: ... def corr( self, - other: DataFrame | Series | np.ndarray | None = ..., + other: DataFrame | Series | None = ..., pairwise: bool | None = ..., - **kwargs, - ) -> DataFrame | Series: ... + ddof: int = ..., + ) -> NDFrameT: ... -class Rolling(_Rolling_and_Expanding): - def is_datetimelike(self) -> bool: ... - win_freq = ... - window = ... - win_type: str = ... - min_periods: int = ... - def count(self) -> DataFrame | Series: ... +class Rolling(RollingAndExpandingMixin[NDFrameT]): + @overload + def aggregate( + self: Rolling[Series], func: AggFuncTypeBase, *args: Any, **kwargs: Any + ) -> Series: ... + @overload + def aggregate( + self: Rolling[Series], + func: AggFuncTypeSeriesToFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... + @overload + def aggregate( + self: Rolling[DataFrame], + func: AggFuncTypeFrame, + *args: Any, + **kwargs: Any, + ) -> DataFrame: ... def apply( self, - func, + func: Callable[..., Any], raw: bool = ..., - engine: str = ..., - engine_kwargs=..., - args=..., - kwargs=..., - ): ... + engine: WindowingEngine = ..., + engine_kwargs: WindowingEngineKwargs | None = ..., + args: tuple[Any, ...] | None = ..., + kwargs: dict[str, Any] | None = ..., + ) -> NDFrameT: ... -class RollingGroupby(Rolling): ... +class RollingGroupby(BaseWindowGroupby[NDFrameT], Rolling): ... diff --git a/pandas-stubs/io/formats/style.pyi b/pandas-stubs/io/formats/style.pyi index df8f9d4f6..f2bad89ce 100644 --- a/pandas-stubs/io/formats/style.pyi +++ b/pandas-stubs/io/formats/style.pyi @@ -18,6 +18,7 @@ from pandas._typing import ( IndexLabel, IntervalClosedType, Level, + QuantileInterpolation, Scalar, T, WriteBuffer, @@ -312,9 +313,7 @@ class Styler(StylerRenderer[Styler]): axis: AxisType | None = ..., q_left: float = ..., q_right: float = ..., - interpolation: Literal[ - "linear", "lower", "higher", "midpoint", "nearest" - ] = ..., + interpolation: QuantileInterpolation = ..., inclusive: IntervalClosedType = ..., props: str | None = ..., ) -> Styler: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index c1e66c3fe..7368b1435 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -736,49 +736,35 @@ def test_types_window() -> None: df.rolling(2, axis=1, center=True) check( - assert_type(df.rolling(2).agg("max"), Union[Scalar, pd.DataFrame, pd.Series]), + assert_type(df.rolling(2).agg("max"), pd.DataFrame), pd.DataFrame, ) check( - assert_type(df.rolling(2).agg(max), Union[Scalar, pd.DataFrame, pd.Series]), + assert_type(df.rolling(2).agg(max), pd.DataFrame), pd.DataFrame, ) check( - assert_type( - df.rolling(2).agg(["max", "min"]), Union[Scalar, pd.DataFrame, pd.Series] - ), + assert_type(df.rolling(2).agg(["max", "min"]), pd.DataFrame), pd.DataFrame, ) check( - assert_type( - df.rolling(2).agg([max, min]), Union[Scalar, pd.DataFrame, pd.Series] - ), + assert_type(df.rolling(2).agg([max, min]), pd.DataFrame), pd.DataFrame, ) check( - assert_type( - df.rolling(2).agg({"col2": "max"}), Union[Scalar, pd.DataFrame, pd.Series] - ), + assert_type(df.rolling(2).agg({"col2": "max"}), pd.DataFrame), pd.DataFrame, ) check( - assert_type( - df.rolling(2).agg({"col2": max}), Union[Scalar, pd.DataFrame, pd.Series] - ), + assert_type(df.rolling(2).agg({"col2": max}), pd.DataFrame), pd.DataFrame, ) check( - assert_type( - df.rolling(2).agg({"col2": ["max", "min"]}), - Union[Scalar, pd.DataFrame, pd.Series], - ), + assert_type(df.rolling(2).agg({"col2": ["max", "min"]}), pd.DataFrame), pd.DataFrame, ) check( - assert_type( - df.rolling(2).agg({"col2": [max, min]}), - Union[Scalar, pd.DataFrame, pd.Series], - ), + assert_type(df.rolling(2).agg({"col2": [max, min]}), pd.DataFrame), pd.DataFrame, ) diff --git a/tests/test_series.py b/tests/test_series.py index 175dd638d..f37ccdb50 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -12,7 +12,6 @@ Iterator, List, Sequence, - Union, cast, ) @@ -517,23 +516,19 @@ def test_types_window() -> None: s.rolling(2, axis=0, center=True) check( - assert_type(s.rolling(2).agg("sum"), Union[Scalar, pd.Series, pd.DataFrame]), + assert_type(s.rolling(2).agg("sum"), pd.Series), pd.Series, ) check( - assert_type(s.rolling(2).agg(sum), Union[Scalar, pd.Series, pd.DataFrame]), + assert_type(s.rolling(2).agg(sum), pd.Series), pd.Series, ) check( - assert_type( - s.rolling(2).agg(["max", "min"]), Union[Scalar, pd.Series, pd.DataFrame] - ), + assert_type(s.rolling(2).agg(["max", "min"]), pd.DataFrame), pd.DataFrame, ) check( - assert_type( - s.rolling(2).agg([max, min]), Union[Scalar, pd.Series, pd.DataFrame] - ), + assert_type(s.rolling(2).agg([max, min]), pd.DataFrame), pd.DataFrame, ) diff --git a/tests/test_windowing.py b/tests/test_windowing.py new file mode 100644 index 000000000..f241bf58a --- /dev/null +++ b/tests/test_windowing.py @@ -0,0 +1,306 @@ +import numpy as np +from pandas import ( + DataFrame, + Series, + date_range, +) +from pandas.core.window import ( + Rolling, + Window, +) +from typing_extensions import assert_type + +from tests import check + +IDX = date_range("1/1/2000", periods=700, freq="D") +S = Series(np.random.standard_normal(700)) +DF = DataFrame({"col1": S, "col2": S}) + + +def test_rolling_basic() -> None: + check(assert_type(DF.rolling(10, win_type="gaussian"), "Window[DataFrame]"), Window) + check(assert_type(DF.rolling(10, min_periods=10), "Rolling[DataFrame]"), Rolling) + + +def test_rolling_basic_math() -> None: + check(assert_type(DF.rolling(10, min_periods=10).count(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).sum(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).mean(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).median(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).var(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).std(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).min(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).max(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).corr(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).cov(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).skew(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).kurt(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10, min_periods=10).sem(), DataFrame), DataFrame) + check(assert_type(DF.rolling(10).quantile(0.5), DataFrame), DataFrame) + check(assert_type(DF.rolling(10).rank("average"), DataFrame), DataFrame) + check(assert_type(DF.rolling(10).rank("min"), DataFrame), DataFrame) + check(assert_type(DF.rolling(10).rank("max"), DataFrame), DataFrame) + + +def test_rolling_apply() -> None: + check(assert_type(DF.rolling(10).apply(np.mean), DataFrame), DataFrame) + + def _mean(df: DataFrame) -> Series: + return df.mean() + + check(assert_type(DF.rolling(10).apply(_mean), DataFrame), DataFrame) + + def _mean2(df: DataFrame) -> np.ndarray: + return np.mean(df, axis=0) + + check(assert_type(DF.rolling(10).apply(_mean2, raw=True), DataFrame), DataFrame) + + def _mean4(df: DataFrame) -> float: + return float(np.mean(df)) + + check(assert_type(DF.rolling(10).apply(_mean4, raw=True), DataFrame), DataFrame) + + +def test_rolling_aggregate() -> None: + check(assert_type(DF.rolling(10).aggregate(np.mean), DataFrame), DataFrame) + check( + assert_type(DF.rolling(10).aggregate(["mean", np.mean]), DataFrame), DataFrame + ) + check( + assert_type( + DF.rolling(10).aggregate({"col1": "mean", "col2": np.mean}), DataFrame + ), + DataFrame, + ) + check(assert_type(DF.rolling(10).agg("sum"), DataFrame), DataFrame) + + check(assert_type(DF.rolling(10).aggregate(np.mean), DataFrame), DataFrame) + check(assert_type(DF.rolling(10).aggregate("mean"), DataFrame), DataFrame) + + def _mean(df: DataFrame) -> Series: + return df.mean() + + check(assert_type(DF.rolling(10).aggregate(_mean), DataFrame), DataFrame) + + check(assert_type(DF.rolling(10).aggregate([np.mean]), DataFrame), DataFrame) + check( + assert_type(DF.rolling(10).aggregate([np.mean, "mean"]), DataFrame), DataFrame + ) + check( + assert_type( + DF.rolling(10).aggregate({"col1": np.mean, "col2": "mean"}), DataFrame + ), + DataFrame, + ) + check( + assert_type( + DF.rolling(10).aggregate({"col1": [np.mean, "mean"], "col2": "mean"}), + DataFrame, + ), + DataFrame, + ) + + # func: np.ufunc | Callable | str | list[Callable | str, np.ufunc] | dict[Hashable, Callable | str | np.ufunc| list[Callable | str]] + check(assert_type(DF.rolling(10).agg("sum"), DataFrame), DataFrame) + + +def test_rolling_basic_math_series() -> None: + check(assert_type(S.rolling(10, min_periods=10).count(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).sum(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).mean(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).median(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).var(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).std(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).min(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).max(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).corr(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).cov(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).skew(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).kurt(), Series), Series) + check(assert_type(S.rolling(10, min_periods=10).sem(), Series), Series) + check(assert_type(S.rolling(10).quantile(0.5), Series), Series) + check(assert_type(S.rolling(10).rank("average"), Series), Series) + check(assert_type(S.rolling(10).rank("min"), Series), Series) + check(assert_type(S.rolling(10).rank("max"), Series), Series) + + +def test_rolling_apply_series() -> None: + check(assert_type(S.rolling(10).apply(np.mean), Series), Series) + + def _mean(df: Series) -> float: + return df.mean() + + check(assert_type(S.rolling(10).apply(_mean), Series), Series) + + def _mean2(df: Series) -> np.ndarray: + return np.mean(df, axis=0) + + check(assert_type(S.rolling(10).apply(_mean2, raw=True), Series), Series) + + +def test_rolling_aggregate_series() -> None: + check(assert_type(S.rolling(10).aggregate(np.mean), Series), Series) + check(assert_type(S.rolling(10).aggregate("mean"), Series), Series) + + def _mean(s: Series) -> float: + return s.mean() + + check(assert_type(S.rolling(10).aggregate(_mean), Series), Series) + + check(assert_type(S.rolling(10).aggregate([np.mean]), DataFrame), DataFrame) + check(assert_type(S.rolling(10).aggregate([np.mean, "mean"]), DataFrame), DataFrame) + check( + assert_type( + S.rolling(10).aggregate({"col1": np.mean, "col2": "mean", "col3": _mean}), + DataFrame, + ), + DataFrame, + ) + check(assert_type(S.rolling(10).agg("sum"), Series), Series) + + +def test_expanding_basic_math() -> None: + check(assert_type(DF.expanding(10).count(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).sum(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).mean(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).median(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).var(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).std(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).min(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).max(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).corr(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).cov(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).skew(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).kurt(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).sem(), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).quantile(0.5), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).rank("average"), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).rank("min"), DataFrame), DataFrame) + check(assert_type(DF.expanding(10).rank("max"), DataFrame), DataFrame) + + +def test_expanding_apply() -> None: + check(assert_type(DF.expanding(10).apply(np.mean), DataFrame), DataFrame) + + def _mean(df: DataFrame) -> Series: + return df.mean() + + check(assert_type(DF.expanding(10).apply(_mean), DataFrame), DataFrame) + + def _mean2(df: DataFrame) -> np.ndarray: + return np.mean(df, axis=0) + + check(assert_type(DF.expanding(10).apply(_mean2, raw=True), DataFrame), DataFrame) + + def _mean4(df: DataFrame) -> float: + return float(np.mean(df)) + + check(assert_type(DF.expanding(10).apply(_mean4, raw=True), DataFrame), DataFrame) + + +def test_expanding_aggregate() -> None: + check(assert_type(DF.expanding(10).aggregate(np.mean), DataFrame), DataFrame) + check( + assert_type(DF.expanding(10).aggregate(["mean", np.mean]), DataFrame), DataFrame + ) + check( + assert_type( + DF.expanding(10).aggregate({"col1": "mean", "col2": np.mean}), DataFrame + ), + DataFrame, + ) + check(assert_type(DF.expanding(10).agg("sum"), DataFrame), DataFrame) + + +def test_expanding_basic_math_series() -> None: + check(assert_type(S.expanding(10).count(), Series), Series) + check(assert_type(S.expanding(10).sum(), Series), Series) + check(assert_type(S.expanding(10).mean(), Series), Series) + check(assert_type(S.expanding(10).median(), Series), Series) + check(assert_type(S.expanding(10).var(), Series), Series) + check(assert_type(S.expanding(10).std(), Series), Series) + check(assert_type(S.expanding(10).min(), Series), Series) + check(assert_type(S.expanding(10).max(), Series), Series) + check(assert_type(S.expanding(10).corr(), Series), Series) + check(assert_type(S.expanding(10).cov(), Series), Series) + check(assert_type(S.expanding(10).skew(), Series), Series) + check(assert_type(S.expanding(10).kurt(), Series), Series) + check(assert_type(S.expanding(10).sem(), Series), Series) + check(assert_type(S.expanding(10).quantile(0.5), Series), Series) + check(assert_type(S.expanding(10).rank("average"), Series), Series) + check(assert_type(S.expanding(10).rank("min"), Series), Series) + check(assert_type(S.expanding(10).rank("max"), Series), Series) + + +def test_expanding_apply_series() -> None: + check(assert_type(S.expanding(10).apply(np.mean), Series), Series) + + def _mean(df: Series) -> float: + return df.mean() + + check(assert_type(S.expanding(10).apply(_mean), Series), Series) + + def _mean2(df: Series) -> np.ndarray: + return np.mean(df, axis=0) + + check(assert_type(S.expanding(10).apply(_mean2, raw=True), Series), Series) + + +def test_expanding_aggregate_series() -> None: + check(assert_type(S.expanding(10).aggregate(np.mean), Series), Series) + check( + assert_type(S.expanding(10).aggregate(["mean", np.mean]), DataFrame), DataFrame + ) + check( + assert_type( + S.expanding(10).aggregate({"col1": "mean", "col2": np.mean}), DataFrame + ), + DataFrame, + ) + check(assert_type(S.expanding(10).agg("sum"), Series), Series) + + +def test_ewm_basic_math() -> None: + check(assert_type(DF.ewm(span=10).sum(), DataFrame), DataFrame) + check(assert_type(DF.ewm(span=10).mean(), DataFrame), DataFrame) + check(assert_type(DF.ewm(span=10).var(), DataFrame), DataFrame) + check(assert_type(DF.ewm(span=10).std(), DataFrame), DataFrame) + check(assert_type(DF.ewm(span=10).corr(), DataFrame), DataFrame) + check(assert_type(DF.ewm(span=10).cov(), DataFrame), DataFrame) + + +def test_ewm_aggregate() -> None: + check(assert_type(DF.ewm(span=10).aggregate(np.mean), DataFrame), DataFrame) + check( + assert_type(DF.ewm(span=10).aggregate(["mean", np.mean]), DataFrame), DataFrame + ) + check( + assert_type( + DF.ewm(span=10).aggregate({"col1": "mean", "col2": np.mean}), DataFrame + ), + DataFrame, + ) + check(assert_type(DF.ewm(span=10).agg("sum"), DataFrame), DataFrame) + + +def test_ewm_basic_math_series() -> None: + check(assert_type(S.ewm(span=10).sum(), Series), Series) + check(assert_type(S.ewm(span=10).mean(), Series), Series) + check(assert_type(S.ewm(span=10).var(), Series), Series) + check(assert_type(S.ewm(span=10).std(), Series), Series) + check(assert_type(S.ewm(span=10).corr(), Series), Series) + check(assert_type(S.ewm(span=10).cov(), Series), Series) + + +def test_ewm_aggregate_series() -> None: + check(assert_type(S.ewm(span=10).aggregate(np.mean), Series), Series) + check( + assert_type(S.ewm(span=10).aggregate(["mean", np.mean]), DataFrame), DataFrame + ) + check( + assert_type( + S.ewm(span=10).aggregate({"col1": "mean", "col2": np.mean}), DataFrame + ), + DataFrame, + ) + check(assert_type(S.ewm(span=10).agg("sum"), Series), Series)