Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Oct 15, 2020
11 changes: 5 additions & 6 deletions pandas/_libs/window/aggregations.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
19 changes: 16 additions & 3 deletions pandas/core/window/ewm.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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
----------
Expand Down Expand Up @@ -60,6 +64,15 @@ def get_center_of_mass(
return float(comass)


def wrap_result(obj: "Series", result: np.ndarray) -> "Series":
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

for clarify should make this wrap_result_1d, but can be in the future

"""
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.
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
36 changes: 26 additions & 10 deletions pandas/core/window/expanding.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
from textwrap import dedent
from typing import Dict, Optional
from typing import 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

Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -135,12 +140,12 @@ def count(self, **kwargs):
@Appender(_shared_docs["apply"])
def apply(
self,
func,
func: Callable,
raw: bool = False,
engine: Optional[str] = None,
engine_kwargs: Optional[Dict[str, bool]] = None,
args=None,
kwargs=None,
args: Optional[Tuple] = None,
kwargs: Optional[Dict] = None,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we will eventually want to use disallow_any_generic as we increase strictness checks xref #30539

so for kwargs could probably use Dict[str, Any] and for args, Tuple[Any, ...]

can you also add type parameters to Callable.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good. Done

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @mroeschke lgtm

):
return super().apply(
func,
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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)


Expand Down
Loading