diff --git a/saqc/lib/rolling.py b/saqc/lib/rolling.py
index 3266cbafeb04884d7bad0336e1222fab758f3da6..49e7f431654cef5446e1d63335d5590feaabf889 100644
--- a/saqc/lib/rolling.py
+++ b/saqc/lib/rolling.py
@@ -105,8 +105,8 @@ class _CustomBaseIndexer(BaseIndexer):
if center:
# centering of dtlike windows is just looking left and right
# with half amount of window-size
- ws1 = (winsz + 1) // 2
- ws2 = winsz // 2
+ ws2, ws1 = divmod(winsz, 2)
+ ws1 += ws2
if self.forward:
ws1, ws2 = ws2, ws1
return ws1, ws2
@@ -382,9 +382,10 @@ def customRoller(obj, window, min_periods=None, # aka minimum non-nan values
def new_count():
self = roller
if not x.is_freq_type:
- result = obj.notna().astype(int)
- theirs.update(min_periods=min_periods or 0)
- return customRoller(result, window, **theirs, **ours).sum()
+ obj_new = obj.notna().astype(int)
+ if min_periods is None:
+ theirs.update(min_periods=0)
+ return obj_new.rolling(indexer, center=None, **theirs).sum()
return self._old_count()
roller._old_count = roller.count
diff --git a/test/lib/test_rolling.py b/test/lib/test_rolling.py
index d8365f4eb9d1decf03cc157b5dfa4e731d8d20f3..b7245b3b0c7564c8c8ba50a8f643597cd84487d0 100644
--- a/test/lib/test_rolling.py
+++ b/test/lib/test_rolling.py
@@ -26,23 +26,14 @@ def data_():
return s
-def data__():
- s1 = pd.Series(1., index=pd.date_range("1999/12", periods=4, freq='1M') + pd.Timedelta('1d'))
- s2 = pd.Series(1., index=pd.date_range('2000/05/15', periods=2, freq='1d'))
- s = pd.concat([s1, s2]).sort_index()
- s.name = 's'
- s[5] = np.nan
- return s
-
-
len_s = len(data_())
def make_num_kws():
l = []
- n = [0, 1, 2, 10, 20, 30]
- mp = [0, 1, 2, 10, 20, 30]
+ n = list(range(len_s))
for window in n:
+ mp = list(range(window))
for min_periods in [None] + mp:
if min_periods is not None and min_periods > window:
continue
@@ -54,7 +45,7 @@ def make_num_kws():
def make_dt_kws():
l = []
n = [0, 1, 2, 10, 32, 70, 120]
- mp = [0, 1, 2, 10, 20, 30]
+ mp = list(range(len_s))
for closed in ['right', 'both', 'neither', 'left']:
for window in n:
for min_periods in [None] + mp:
@@ -80,74 +71,115 @@ def print_diff(s, result, expected):
print(df)
-def runtest_for_kw_combi(s, kws, func='sum'):
- print(kws)
- forward = kws.pop('forward', False)
-
- def calc(roller):
- if isinstance(func, str):
- return getattr(roller, func)()
- else:
- return getattr(roller, 'apply')(func)
-
- if forward:
-
- try:
- expR = pd.Series(reversed(s), reversed(s.index)).rolling(**kws)
- expected = calc(expR)[::-1]
- except Exception:
- pytest.skip("pandas failed")
- return
-
- resR = customRoller(s, forward=True, **kws)
- result = calc(resR)
-
- success = check_series(result, expected)
- if not success:
- print_diff(s, result, expected)
- assert False, f"forward=True !! {kws}"
+def call_rolling_function(roller, func):
+ if isinstance(func, str):
+ return getattr(roller, func)()
else:
-
- try:
- expR = s.rolling(**kws)
- expected = calc(expR)
- except Exception:
- pytest.skip("pandas failed")
- return
-
- resR = customRoller(s, **kws)
- result = calc(resR)
-
- success = check_series(result, expected)
- if not success:
- print_diff(s, result, expected)
- assert False
+ return getattr(roller, 'apply')(func)
@pytest.mark.parametrize("kws", make_dt_kws(), ids=lambda x: str(x))
@pytest.mark.parametrize("func", FUNCTS)
def test_pandas_conform_dt(data, kws, func):
- runtest_for_kw_combi(data, kws, func=func)
+ s = data
+ try:
+ expR = s.rolling(**kws)
+ expected = call_rolling_function(expR, func)
+ except Exception as e0:
+ # pandas failed, so we should also fail
+ try:
+ resR = customRoller(s, **kws)
+ result = call_rolling_function(resR, func)
+ except Exception as e1:
+ assert type(e0) == type(e1)
+ return
+ assert False, 'pandas faild, but we succeed'
+
+ resR = customRoller(s, **kws)
+ result = call_rolling_function(resR, func)
+ success = check_series(result, expected)
+ if success:
+ return
+ print_diff(s, result, expected)
+ assert False
@pytest.mark.parametrize("kws", make_num_kws(), ids=lambda x: str(x))
@pytest.mark.parametrize("func", FUNCTS)
def test_pandas_conform_num(data, kws, func):
- runtest_for_kw_combi(data, kws, func=func)
+ s = data
+ try:
+ expR = s.rolling(**kws)
+ expected = call_rolling_function(expR, func)
+ except Exception as e0:
+ # pandas failed, so we should also fail
+ try:
+ resR = customRoller(s, **kws)
+ result = call_rolling_function(resR, func)
+ except Exception as e1:
+ assert type(e0) == type(e1)
+ return
+ assert False, 'pandas faild, but we succeed'
+
+ resR = customRoller(s, **kws)
+ result = call_rolling_function(resR, func)
+ success = check_series(result, expected)
+ if success:
+ return
+ print_diff(s, result, expected)
+ assert False
@pytest.mark.parametrize("kws", make_dt_kws(), ids=lambda x: str(x))
@pytest.mark.parametrize("func", FUNCTS)
def test_forward_dt(data, kws, func):
- kws.update(forward=True)
- runtest_for_kw_combi(data, kws, func=func)
+ s = data
+ try:
+ expR = pd.Series(reversed(s), reversed(s.index)).rolling(**kws)
+ expected = call_rolling_function(expR, func)[::-1]
+ except Exception as e0:
+ # pandas failed, so we should also fail
+ try:
+ resR = customRoller(s, forward=True, **kws)
+ result = call_rolling_function(resR, func)
+ except Exception as e1:
+ assert type(e0) == type(e1)
+ return
+ assert False, 'pandas faild, but we succeed'
+
+ resR = customRoller(s, forward=True, **kws)
+ result = call_rolling_function(resR, func)
+ success = check_series(result, expected)
+ if success:
+ return
+ print_diff(s, result, expected)
+ assert False
@pytest.mark.parametrize("kws", make_num_kws(), ids=lambda x: str(x))
@pytest.mark.parametrize("func", FUNCTS)
def test_forward_num(data, kws, func):
- kws.update(forward=True)
- runtest_for_kw_combi(data, kws, func=func)
+ s = data
+ try:
+ expR = pd.Series(reversed(s), reversed(s.index)).rolling(**kws)
+ expected = call_rolling_function(expR, func)[::-1]
+ except Exception as e0:
+ # pandas failed, so we should also fail
+ try:
+ resR = customRoller(s, forward=True, **kws)
+ result = call_rolling_function(resR, func)
+ except Exception as e1:
+ assert type(e0) == type(e1)
+ return
+ assert False, 'pandas faild, but we succeed'
+
+ resR = customRoller(s, forward=True, **kws)
+ result = call_rolling_function(resR, func)
+ success = check_series(result, expected)
+ if success:
+ return
+ print_diff(s, result, expected)
+ assert False
def dt_center_kws():