clean up

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

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():