8000 TYPING: pandas/core/window.py (#27391) · pandas-dev/pandas@a4c19e7 · GitHub
[go: up one dir, main page]
Skip to content

Commit a4c19e7

Browse files
mroeschkemroeschke
authored
TYPING: pandas/core/window.py (#27391)
* TYPING: pandas/core/window.py * Address comments * Remove unused import * Sort imports
1 parent 479d003 commit a4c19e7

File tree

2 files changed

+63
-57
lines changed

2 files changed

+63
-57
lines changed

‎pandas/_typing.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,4 @@
3131

3232
FrameOrSeries = TypeVar("FrameOrSeries", ABCSeries, ABCDataFrame)
3333
Scalar = Union[str, int, float]
34+
Axis = Union[str, int]

‎pandas/core/window.py

Lines changed: 62 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from collections import defaultdict
66
from datetime import timedelta
77
from textwrap import dedent
8-
from typing import Set
8+
from typing import List, Optional, Set
99
import warnings
1010

1111
import numpy as np
@@ -35,6 +35,7 @@
3535
ABCTimedeltaIndex,
3636
)
3737

38+
from pandas._typing import Axis, FrameOrSeries
3839
from pandas.core.base import DataError, PandasObject, SelectionMixin
3940
import pandas.core.common as com
4041
from pandas.core.generic import _shared_docs
@@ -63,24 +64,23 @@ class _Window(PandasObject, SelectionMixin):
6364
"axis",
6465
"on",
6566
"closed",
66-
]
67+
] # type: List[str]
6768
exclusions = set() # type: Set[str]
6869

6970
def __init__(
7071
self,
7172
obj,
7273
window=None,
73-
min_periods=None,
74-
center=False,
75-
win_type=None,
76-
axis=0,
77-
on=None,
78-
closed=None,
74+
min_periods: Optional[int] = None,
75+
center: Optional[bool] = False,
76+
win_type: Optional[str] = None,
77+
axis: Axis = 0,
78+
on: Optional[str] = None,
79+
closed: Optional[str] = None,
7980
**kwargs
8081
):
8182

8283
self.__dict__.update(kwargs)
83-
self.blocks = []
8484
self.obj = obj
8585
self.on = on
8686
self.closed = closed
@@ -97,15 +97,15 @@ def _constructor(self):
9797
return Window
9898

9999
@property
100-
def is_datetimelike(self):
100+
def is_datetimelike(self) -> Optional[bool]:
101101
return None
102102

103103
@property
104104
def _on(self):
105105
return None
106106

107107
@property
108-
def is_freq_type(self):
108+
def is_freq_type(self) -> bool:
109109
return self.win_type == "freq"
110110

111111
def validate(self):
@@ -121,30 +121,20 @@ def validate(self):
121121
]:
122122
raise ValueError("closed must be 'right', 'left', 'both' or " "'neither'")
123123

124-
def _convert_freq(self):
125-
"""
126-
Resample according to the how, return a new object.
127-
"""
128-
obj = self._selected_obj
129-
index = None
130-
return obj, index
131-
132124
def _create_blocks(self):
133125
"""
134126
Split data into blocks & return conformed data.
135127
"""
136128

137-
obj, index = self._convert_freq()
138-
if index is not None:
139-
index = self._on
129+
obj = self._selected_obj
140130

141131
# filter out the on from the object
142132
if self.on is not None:
143133
if obj.ndim == 2:
144134
obj = obj.reindex(columns=obj.columns.difference([self.on]), copy=False)
145135
blocks = obj._to_dict_of_blocks(copy=False).values()
146136

147-
return blocks, obj, index
137+
return blocks, obj
148138

149139
def _gotitem(self, key, ndim, subset=None):
150140
"""
@@ -186,10 +176,10 @@ def _get_window(self, other=None):
186176
return self.window
187177

188178
@property
189-
def _window_type(self):
179+
def _window_type(self) -> str:
190180
return self.__class__.__name__
191181

192-
def __repr__(self):
182+
def __repr__(self) -> str:
193183
"""
194184
Provide a nice str repr of our rolling object.
195185
"""
@@ -207,23 +197,21 @@ def __iter__(self):
207197
url = "https://github.com/pandas-dev/pandas/issues/11704"
208198
raise NotImplementedError("See issue #11704 {url}".format(url=url))
209199

210-
def _get_index(self, index=None):
200+
def _get_index(self) -> Optional[np.ndarray]:
211201
"""
212-
Return index as ndarrays.
202+
Return index as an ndarray.
213203
214204
Returns
215205
-------
216-
tuple of (index, index_as_ndarray)
206+
None or ndarray
217207
"""
218208

219209
if self.is_freq_type:
220-
if index is None:
221-
index = self._on
222-
return index, index.asi8
223-
return index, index
224-
225-
def _prep_values(self, values=None, kill_inf=True):
210+
return self._on.asi8
211+
return None
226212

213+
def _prep_values(self, values: Optional[np.ndarray] = None) -> np.ndarray:
214+
"""Convert input to numpy arrays for Cython routines"""
227215
if values is None:
228216
values = getattr(self._selected_obj, "values", self._selected_obj)
229217

@@ -247,13 +235,12 @@ def _prep_values(self, values=None, kill_inf=True):
247235
"cannot handle this type -> {0}" "".format(values.dtype)
248236
)
249237

250-
if kill_inf:
251-
values = values.copy()
252-
values[np.isinf(values)] = np.NaN
238+
# Always convert inf to nan
239+
values[np.isinf(values)] = np.NaN
253240

254241
return values
255242

256-
def _wrap_result(self, result, block=None, obj=None):
243+
def _wrap_result(self, result, block=None, obj=None) -> FrameOrSeries:
257244
"""
258245
Wrap a single result.
259246
"""
@@ -281,7 +268,7 @@ def _wrap_result(self, result, block=None, obj=None):
281268
return type(obj)(result, index=index, columns=block.columns)
282269
return result
283270

284-
def _wrap_results(self, results, blocks, obj, exclude=None):
271+
def _wrap_results(self, results, blocks, obj, exclude=None) -> FrameOrSeries:
285272
"""
286273
Wrap the results.
287274
@@ -335,7 +322,7 @@ def _wrap_results(self, results, blocks, obj, exclude=None):
335322
return obj.astype("float64")
336323
return concat(final, axis=1).reindex(columns=columns, copy=False)
337324

338-
def _center_window(self, result, window):
325+
def _center_window(self, result, window) -> np.ndarray:
339326
"""
340327
Center the result in the window.
341328
"""
@@ -724,7 +711,7 @@ def _apply_window(self, mean=True, **kwargs):
724711
window = self._prep_window(**kwargs)
725712
center = self.center
726713

727-
blocks, obj, index = self._create_blocks()
714+
blocks, obj = se F42D lf._create_blocks()
728715
block_list = list(blocks)
729716

730717
results = []
@@ -912,9 +899,9 @@ def _apply(
912899
if check_minp is None:
913900
check_minp = _use_window
914901

915-
blocks, obj, index = self._create_blocks()
902+
blocks, obj = self._create_blocks()
916903
block_list = list(blocks)
917-
index, indexi = self._get_index(index=index)
904+
index_as_array = self._get_index()
918905

919906
results = []
920907
exclude = []
@@ -947,7 +934,7 @@ def func(arg, window, min_periods=None, closed=None):
947934
minp = check_minp(min_periods, window)
948935
# ensure we are only rolling on floats
949936
arg = ensure_float64(arg)
950-
return cfunc(arg, window, minp, indexi, closed, **kwargs)
937+
return cfunc(arg, window, minp, index_as_array, closed, **kwargs)
951938

952939
# calculation function
953940
if center:
@@ -1027,9 +1014,9 @@ class _Rolling_and_Expanding(_Rolling):
10271014

10281015
def count(self):
10291016

1030-
blocks, obj, index = self._create_blocks()
1017+
blocks, obj = self._create_blocks()
10311018
# Validate the index
1032-
self._get_index(index=index)
1019+
self._get_index()
10331020

10341021
window = self._get_window()
10351022
window = min(window, len(obj)) if not self.center else window
@@ -1088,11 +1075,10 @@ def count(self):
10881075
def apply(self, func, raw=None, args=(), kwargs={}):
10891076
from pandas import Series
10901077

1091-
# TODO: _level is unused?
1092-
_level = kwargs.pop("_level", None) # noqa
1078+
kwargs.pop("_level", None)
10931079
window = self._get_window()
10941080< F438 /td>
offset = _offset(window, self.center)
1095-
index, indexi = self._get_index()
1081+
index_as_array = self._get_index()
10961082

10971083
# TODO: default is for backward compat
10981084
# change to False in the future
@@ -1113,7 +1099,16 @@ def f(arg, window, min_periods, closed):
11131099
if not raw:
11141100
arg = Series(arg, index=self.obj.index)
11151101
return libwindow.roll_generic(
1116-
arg, window, minp, indexi, closed, offset, func, raw, args, kwargs
1102+
arg,
1103+
window,
1104+
minp,
1105+
index_as_array,
1106+
closed,
1107+
offset,
1108+
func,
1109+
raw,
1110+
args,
1111+
kwargs,
11171112
)
11181113

11191114
return self._apply(f, func, args=args, kwargs=kwargs, center=False, raw=raw)
@@ -1285,12 +1280,12 @@ def median(self, **kwargs):
12851280
def std(self, ddof=1, *args, **kwargs):
12861281
nv.validate_window_func("std", args, kwargs)
12871282
window = self._get_window()
1288-
index, indexi = self._get_index()
1283+
index_as_array = self._get_index()
12891284

12901285
def f(arg, *args, **kwargs):
12911286
minp = _require_min_periods(1)(self.min_periods, window)
12921287
return _zsqrt(
1293-
libwindow.roll_var(arg, window, minp, indexi, self.closed, ddof)
1288+
libwindow.roll_var(arg, window, minp, index_as_array, self.closed, ddof)
12941289
)
12951290

12961291
return self._apply(
@@ -1474,17 +1469,27 @@ def kurt(self, **kwargs):
14741469

14751470
def quantile(self, quantile, interpo 10000 lation="linear", **kwargs):
14761471
window = self._get_window()
1477-
index, indexi = self._get_index()
1472+
index_as_array = self._get_index()
14781473

14791474
def f(arg, *args, **kwargs):
14801475
minp = _use_window(self.min_periods, window)
14811476
if quantile == 1.0:
1482-
return libwindow.roll_max(arg, window, minp, indexi, self.closed)
1477+
return libwindow.roll_max(
1478+
arg, window, minp, index_as_array, self.closed
1479+
)
14831480
elif quantile == 0.0:
1484-
return libwindow.roll_min(arg, window, minp, indexi, self.closed)
1481+
return libwindow.roll_min(
1482+
arg, window, minp, index_as_array, self.closed
1483+
)
14851484
else:
14861485
return libwindow.roll_quantile(
1487-
arg, window, minp, indexi, self.closed, quantile, interpolation
1486+
arg,
1487+
window,
1488+
minp,
1489+
index_as_array,
1490+
self.closed,
1491+
quantile,
1492+
interpolation,
14881493
)
14891494

14901495
return self._apply(f, "quantile", quantile=quantile, **kwargs)
@@ -2450,7 +2455,7 @@ def _apply(self, func, **kwargs):
24502455
-------
24512456
y : same type as input argument
24522457
"""
2453-
blocks, obj, index = self._create_blocks()
2458+
blocks, obj = self._create_blocks()
24542459
block_list = list(blocks)
24552460

24562461
results = []

0 commit comments

Comments
 (0)
0