diff --git a/pandas/core/window/ewm.py b/pandas/core/window/ewm.py index 5a71db82f26e4..57fb2381b71c8 100644 --- a/pandas/core/window/ewm.py +++ b/pandas/core/window/ewm.py @@ -14,6 +14,7 @@ from pandas._libs.tslibs import Timedelta import pandas._libs.window.aggregations as window_aggregations from pandas._typing import ( + Axis, FrameOrSeries, FrameOrSeriesUnion, TimedeltaConvertibleTypes, @@ -227,7 +228,7 @@ class ExponentialMovingWindow(BaseWindow): def __init__( self, - obj, + obj: FrameOrSeries, com: Optional[float] = None, span: Optional[float] = None, halflife: Optional[Union[float, TimedeltaConvertibleTypes]] = None, @@ -235,7 +236,7 @@ def __init__( min_periods: int = 0, adjust: bool = True, ignore_na: bool = False, - axis: int = 0, + axis: Axis = 0, times: Optional[Union[str, np.ndarray, FrameOrSeries]] = None, ): super().__init__( diff --git a/pandas/core/window/expanding.py b/pandas/core/window/expanding.py index 5f5b0bf8552ed..8b7182458dd1f 100644 --- a/pandas/core/window/expanding.py +++ b/pandas/core/window/expanding.py @@ -10,7 +10,11 @@ import numpy as np -from pandas._typing import FrameOrSeries +from pandas._typing import ( + Axis, + FrameOrSeries, + FrameOrSeriesUnion, +) from pandas.compat.numpy import function as nv from pandas.util._decorators import doc @@ -94,7 +98,14 @@ class Expanding(RollingAndExpandingMixin): _attributes = ["min_periods", "center", "axis", "method"] - def __init__(self, obj, min_periods=1, center=None, axis=0, method="single"): + def __init__( + self, + obj: FrameOrSeries, + min_periods: int = 1, + center=None, + axis: Axis = 0, + method: str = "single", + ): super().__init__( obj=obj, min_periods=min_periods, center=center, axis=axis, method=method ) @@ -200,7 +211,13 @@ def apply( aggregation_description="sum", agg_method="sum", ) - def sum(self, *args, engine=None, engine_kwargs=None, **kwargs): + def sum( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_expanding_func("sum", args, kwargs) return super().sum(*args, engine=engine, engine_kwargs=engine_kwargs, **kwargs) @@ -220,7 +237,13 @@ def sum(self, *args, engine=None, engine_kwargs=None, **kwargs): aggregation_description="maximum", agg_method="max", ) - def max(self, *args, engine=None, engine_kwargs=None, **kwargs): + def max( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_expanding_func("max", args, kwargs) return super().max(*args, engine=engine, engine_kwargs=engine_kwargs, **kwargs) @@ -240,7 +263,13 @@ def max(self, *args, engine=None, engine_kwargs=None, **kwargs): aggregation_description="minimum", agg_method="min", ) - def min(self, *args, engine=None, engine_kwargs=None, **kwargs): + def min( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_expanding_func("min", args, kwargs) return super().min(*args, engine=engine, engine_kwargs=engine_kwargs, **kwargs) @@ -260,7 +289,13 @@ def min(self, *args, engine=None, engine_kwargs=None, **kwargs): aggregation_description="mean", agg_method="mean", ) - def mean(self, *args, engine=None, engine_kwargs=None, **kwargs): + def mean( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_expanding_func("mean", args, kwargs) return super().mean(*args, engine=engine, engine_kwargs=engine_kwargs, **kwargs) @@ -279,7 +314,12 @@ def mean(self, *args, engine=None, engine_kwargs=None, **kwargs): aggregation_description="median", agg_method="median", ) - def median(self, engine=None, engine_kwargs=None, **kwargs): + def median( + self, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): return super().median(engine=engine, engine_kwargs=engine_kwargs, **kwargs) @doc( @@ -508,8 +548,8 @@ def kurt(self, **kwargs): ) def quantile( self, - quantile, - interpolation="linear", + quantile: float, + interpolation: str = "linear", **kwargs, ): return super().quantile( @@ -549,7 +589,7 @@ def quantile( ) def cov( self, - other: Optional[Union[np.ndarray, FrameOrSeries]] = None, + other: Optional[Union[np.ndarray, FrameOrSeriesUnion]] = None, pairwise: Optional[bool] = None, ddof: int = 1, **kwargs, @@ -614,7 +654,7 @@ def cov( ) def corr( self, - other: Optional[Union[np.ndarray, FrameOrSeries]] = None, + other: Optional[Union[np.ndarray, FrameOrSeriesUnion]] = None, pairwise: Optional[bool] = None, ddof: int = 1, **kwargs, diff --git a/pandas/core/window/rolling.py b/pandas/core/window/rolling.py index 299e1755c0025..e6577b0b715fa 100644 --- a/pandas/core/window/rolling.py +++ b/pandas/core/window/rolling.py @@ -550,7 +550,7 @@ class BaseWindowGroupby(BaseWindow): def __init__( self, - obj, + obj: FrameOrSeries, *args, _grouper=None, **kwargs, @@ -1184,7 +1184,13 @@ def apply_func(values, begin, end, min_periods, raw=raw): return apply_func - def sum(self, *args, engine=None, engine_kwargs=None, **kwargs): + def sum( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_window_func("sum", args, kwargs) if maybe_use_numba(engine): if self.method == "table": @@ -1201,7 +1207,13 @@ def sum(self, *args, engine=None, engine_kwargs=None, **kwargs): window_func = window_aggregations.roll_sum return self._apply(window_func, name="sum", **kwargs) - def max(self, *args, engine=None, engine_kwargs=None, **kwargs): + def max( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_window_func("max", args, kwargs) if maybe_use_numba(engine): if self.method == "table": @@ -1218,7 +1230,13 @@ def max(self, *args, engine=None, engine_kwargs=None, **kwargs): window_func = window_aggregations.roll_max return self._apply(window_func, name="max", **kwargs) - def min(self, *args, engine=None, engine_kwargs=None, **kwargs): + def min( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_window_func("min", args, kwargs) if maybe_use_numba(engine): if self.method == "table": @@ -1235,7 +1253,13 @@ def min(self, *args, engine=None, engine_kwargs=None, **kwargs): window_func = window_aggregations.roll_min return self._apply(window_func, name="min", **kwargs) - def mean(self, *args, engine=None, engine_kwargs=None, **kwargs): + def mean( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_window_func("mean", args, kwargs) if maybe_use_numba(engine): if self.method == "table": @@ -1252,7 +1276,12 @@ def mean(self, *args, engine=None, engine_kwargs=None, **kwargs): window_func = window_aggregations.roll_mean return self._apply(window_func, name="mean", **kwargs) - def median(self, engine=None, engine_kwargs=None, **kwargs): + def median( + self, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): if maybe_use_numba(engine): if self.method == "table": func = generate_manual_numpy_nan_agg_with_axis(np.nanmedian) @@ -1323,7 +1352,13 @@ def quantile(self, quantile: float, interpolation: str = "linear", **kwargs): return self._apply(window_func, name="quantile", **kwargs) - def cov(self, other=None, pairwise=None, ddof=1, **kwargs): + def cov( + self, + other: Optional[Union[np.ndarray, FrameOrSeriesUnion]] = None, + pairwise: Optional[bool] = None, + ddof: int = 1, + **kwargs, + ): from pandas import Series def cov_func(x, y): @@ -1355,7 +1390,13 @@ def cov_func(x, y): return self._apply_pairwise(self._selected_obj, other, pairwise, cov_func) - def corr(self, other=None, pairwise=None, ddof=1, **kwargs): + def corr( + self, + other: Optional[Union[np.ndarray, FrameOrSeriesUnion]] = None, + pairwise: Optional[bool] = None, + ddof: int = 1, + **kwargs, + ): from pandas import Series @@ -1574,7 +1615,13 @@ def count(self): agg_method="apply", ) def apply( - self, func, raw=False, engine=None, engine_kwargs=None, args=None, kwargs=None + self, + func: Callable[..., Any], + raw: bool = False, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + args: Optional[Tuple[Any, ...]] = None, + kwargs: Optional[Dict[str, Any]] = None, ): return super().apply( func, @@ -1649,7 +1696,13 @@ def apply( aggregation_description="sum", agg_method="sum", ) - def sum(self, *args, engine=None, engine_kwargs=None, **kwargs): + def sum( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_rolling_func("sum", args, kwargs) return super().sum(*args, engine=engine, engine_kwargs=engine_kwargs, **kwargs) @@ -1669,7 +1722,13 @@ def sum(self, *args, engine=None, engine_kwargs=None, **kwargs): aggregation_description="maximum", agg_method="max", ) - def max(self, *args, engine=None, engine_kwargs=None, **kwargs): + def max( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_rolling_func("max", args, kwargs) return super().max(*args, engine=engine, engine_kwargs=engine_kwargs, **kwargs) @@ -1704,7 +1763,13 @@ def max(self, *args, engine=None, engine_kwargs=None, **kwargs): aggregation_description="minimum", agg_method="min", ) - def min(self, *args, engine=None, engine_kwargs=None, **kwargs): + def min( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_rolling_func("min", args, kwargs) return super().min(*args, engine=engine, engine_kwargs=engine_kwargs, **kwargs) @@ -1746,7 +1811,13 @@ def min(self, *args, engine=None, engine_kwargs=None, **kwargs): aggregation_description="mean", agg_method="mean", ) - def mean(self, *args, engine=None, engine_kwargs=None, **kwargs): + def mean( + self, + *args, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): nv.validate_rolling_func("mean", args, kwargs) return super().mean(*args, engine=engine, engine_kwargs=engine_kwargs, **kwargs) @@ -1780,7 +1851,12 @@ def mean(self, *args, engine=None, engine_kwargs=None, **kwargs): aggregation_description="median", agg_method="median", ) - def median(self, engine=None, engine_kwargs=None, **kwargs): + def median( + self, + engine: Optional[str] = None, + engine_kwargs: Optional[Dict[str, bool]] = None, + **kwargs, + ): return super().median(engine=engine, engine_kwargs=engine_kwargs, **kwargs) @doc( @@ -1828,7 +1904,7 @@ def median(self, engine=None, engine_kwargs=None, **kwargs): aggregation_description="standard deviation", agg_method="std", ) - def std(self, ddof=1, *args, **kwargs): + def std(self, ddof: int = 1, *args, **kwargs): nv.validate_rolling_func("std", args, kwargs) return super().std(ddof=ddof, **kwargs) @@ -1877,7 +1953,7 @@ def std(self, ddof=1, *args, **kwargs): aggregation_description="variance", agg_method="var", ) - def var(self, ddof=1, *args, **kwargs): + def var(self, ddof: int = 1, *args, **kwargs): nv.validate_rolling_func("var", args, kwargs) return super().var(ddof=ddof, **kwargs) @@ -1933,7 +2009,7 @@ def skew(self, **kwargs): aggregation_description="standard error of mean", agg_method="sem", ) - def sem(self, ddof=1, *args, **kwargs): + def sem(self, ddof: int = 1, *args, **kwargs): return self.std(*args, **kwargs) / (self.count() - ddof).pow(0.5) @doc( @@ -2023,7 +2099,7 @@ def kurt(self, **kwargs): aggregation_description="quantile", agg_method="quantile", ) - def quantile(self, quantile, interpolation="linear", **kwargs): + def quantile(self, quantile: float, interpolation: str = "linear", **kwargs): return super().quantile( quantile=quantile, interpolation=interpolation, @@ -2059,7 +2135,13 @@ def quantile(self, quantile, interpolation="linear", **kwargs): aggregation_description="sample covariance", agg_method="cov", ) - def cov(self, other=None, pairwise=None, ddof=1, **kwargs): + def cov( + self, + other: Optional[Union[np.ndarray, FrameOrSeriesUnion]] = None, + pairwise: Optional[bool] = None, + ddof: int = 1, + **kwargs, + ): return super().cov(other=other, pairwise=pairwise, ddof=ddof, **kwargs) @doc( @@ -2178,7 +2260,13 @@ def cov(self, other=None, pairwise=None, ddof=1, **kwargs): aggregation_description="correlation", agg_method="corr", ) - def corr(self, other=None, pairwise=None, ddof=1, **kwargs): + def corr( + self, + other: Optional[Union[np.ndarray, FrameOrSeriesUnion]] = None, + pairwise: Optional[bool] = None, + ddof: int = 1, + **kwargs, + ): return super().corr(other=other, pairwise=pairwise, ddof=ddof, **kwargs)