REF/TYP: pandas/core/window/*.py (#37091)

Author: mroeschke

pandas/_libs/window/aggregations.pyx

@@ -3,7 +3,6 @@
 import cython
 from cython import Py_ssize_t
 
-from libc.stdlib cimport free, malloc
 from libcpp.deque cimport deque
 
 import numpy as np
@@ -906,7 +905,7 @@ interpolation_types = {
 
 
 def roll_quantile(ndarray[float64_t, cast=True] values, ndarray[int64_t] start,
-                  ndarray[int64_t] end, int64_t minp, int64_t win,
+                  ndarray[int64_t] end, int64_t minp,
                   float64_t quantile, str interpolation):
     """
     O(N log(window)) implementation using skip list
@@ -933,7 +932,7 @@ def roll_quantile(ndarray[float64_t, cast=True] values, ndarray[int64_t] start,
     # actual skiplist ops outweigh any window computation costs
     output = np.empty(N, dtype=float)
 
-    if win == 0 or (end - start).max() == 0:
+    if (end - start).max() == 0:
         output[:] = NaN
         return output
     win = (end - start).max()
@@ -1020,7 +1019,7 @@ def roll_quantile(ndarray[float64_t, cast=True] values, ndarray[int64_t] start,
 def roll_apply(object obj,
                ndarray[int64_t] start, ndarray[int64_t] end,
                int64_t minp,
-               object func, bint raw,
+               object function, bint raw,
                tuple args, dict kwargs):
     cdef:
         ndarray[float64_t] output, counts
@@ -1048,9 +1047,9 @@ def roll_apply(object obj,
 
         if counts[i] >= minp:
             if raw:
-                output[i] = func(arr[s:e], *args, **kwargs)
+                output[i] = function(arr[s:e], *args, **kwargs)
             else:
-                output[i] = func(obj.iloc[s:e], *args, **kwargs)
+                output[i] = function(obj.iloc[s:e], *args, **kwargs)
         else:
             output[i] = NaN
 

pandas/core/window/ewm.py

@@ -1,7 +1,7 @@
 import datetime
 from functools import partial
 from textwrap import dedent
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
 
@@ -17,6 +17,10 @@
 from pandas.core.window.common import _doc_template, _shared_docs, zsqrt
 from pandas.core.window.rolling import BaseWindow, flex_binary_moment
 
+if TYPE_CHECKING:
+    from pandas import Series
+
+
 _bias_template = """
         Parameters
         ----------
@@ -60,6 +64,15 @@ def get_center_of_mass(
     return float(comass)
 
 
+def wrap_result(obj: "Series", result: np.ndarray) -> "Series":
+    """
+    Wrap a single 1D result.
+    """
+    obj = obj._selected_obj
+
+    return obj._constructor(result, obj.index, name=obj.name)
+
+
 class ExponentialMovingWindow(BaseWindow):
     r"""
     Provide exponential weighted (EW) functions.
@@ -413,7 +426,7 @@ def _get_cov(X, Y):
                 self.min_periods,
                 bias,
             )
-            return X._wrap_result(cov)
+            return wrap_result(X, cov)
 
         return flex_binary_moment(
             self._selected_obj, other._selected_obj, _get_cov, pairwise=bool(pairwise)
@@ -467,7 +480,7 @@ def _cov(x, y):
                 x_var = _cov(x_values, x_values)
                 y_var = _cov(y_values, y_values)
                 corr = cov / zsqrt(x_var * y_var)
-            return X._wrap_result(corr)
+            return wrap_result(X, corr)
 
         return flex_binary_moment(
             self._selected_obj, other._selected_obj, _get_corr, pairwise=bool(pairwise)

pandas/core/window/expanding.py

@@ -1,6 +1,9 @@
 from textwrap import dedent
-from typing import Dict, Optional
+from typing import Any, Callable, Dict, Optional, Tuple, Union
 
+import numpy as np
+
+from pandas._typing import FrameOrSeries
 from pandas.compat.numpy import function as nv
 from pandas.util._decorators import Appender, Substitution, doc
 
@@ -65,7 +68,9 @@ def __init__(self, obj, min_periods=1, center=None, axis=0, **kwargs):
     def _constructor(self):
         return Expanding
 
-    def _get_window(self, other=None, **kwargs):
+    def _get_window(
+        self, other: Optional[Union[np.ndarray, FrameOrSeries]] = None, **kwargs
+    ) -> int:
         """
         Get the window length over which to perform some operation.
 
@@ -135,12 +140,12 @@ def count(self, **kwargs):
     @Appender(_shared_docs["apply"])
     def apply(
         self,
-        func,
+        func: Callable[..., Any],
         raw: bool = False,
         engine: Optional[str] = None,
         engine_kwargs: Optional[Dict[str, bool]] = None,
-        args=None,
-        kwargs=None,
+        args: Optional[Tuple[Any, ...]] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
     ):
         return super().apply(
             func,
@@ -183,19 +188,19 @@ def median(self, **kwargs):
 
     @Substitution(name="expanding", versionadded="")
     @Appender(_shared_docs["std"])
-    def std(self, ddof=1, *args, **kwargs):
+    def std(self, ddof: int = 1, *args, **kwargs):
         nv.validate_expanding_func("std", args, kwargs)
         return super().std(ddof=ddof, **kwargs)
 
     @Substitution(name="expanding", versionadded="")
     @Appender(_shared_docs["var"])
-    def var(self, ddof=1, *args, **kwargs):
+    def var(self, ddof: int = 1, *args, **kwargs):
         nv.validate_expanding_func("var", args, kwargs)
         return super().var(ddof=ddof, **kwargs)
 
     @Substitution(name="expanding")
     @Appender(_shared_docs["sem"])
-    def sem(self, ddof=1, *args, **kwargs):
+    def sem(self, ddof: int = 1, *args, **kwargs):
         return super().sem(ddof=ddof, **kwargs)
 
     @Substitution(name="expanding", func_name="skew")
@@ -245,12 +250,23 @@ def quantile(self, quantile, interpolation="linear", **kwargs):
     @Substitution(name="expanding", func_name="cov")
     @Appender(_doc_template)
     @Appender(_shared_docs["cov"])
-    def cov(self, other=None, pairwise=None, ddof=1, **kwargs):
+    def cov(
+        self,
+        other: Optional[Union[np.ndarray, FrameOrSeries]] = None,
+        pairwise: Optional[bool] = None,
+        ddof: int = 1,
+        **kwargs,
+    ):
         return super().cov(other=other, pairwise=pairwise, ddof=ddof, **kwargs)
 
     @Substitution(name="expanding")
     @Appender(_shared_docs["corr"])
-    def corr(self, other=None, pairwise=None, **kwargs):
+    def corr(
+        self,
+        other: Optional[Union[np.ndarray, FrameOrSeries]] = None,
+        pairwise: Optional[bool] = None,
+        **kwargs,
+    ):
         return super().corr(other=other, pairwise=pairwise, **kwargs)