TST: more comprehensive dtype testing for rolling

Author: jreback

pandas/core/window.py

@@ -124,13 +124,17 @@ def _dir_additions(self):
     def _get_window(self, other=None):
         return self.window
 
+    @property
+    def _window_type(self):
+        return self.__class__.__name__
+
     def __unicode__(self):
         """ provide a nice str repr of our rolling object """
 
         attrs = ["{k}={v}".format(k=k, v=getattr(self, k))
                  for k in self._attributes
                  if getattr(self, k, None) is not None]
-        return "{klass} [{attrs}]".format(klass=self.__class__.__name__,
+        return "{klass} [{attrs}]".format(klass=self._window_type,
                                           attrs=','.join(attrs))
 
     def _shallow_copy(self, obj=None, **kwargs):
@@ -155,8 +159,12 @@ def _prep_values(self, values=None, kill_inf=True, how=None):
             values = com._ensure_float64(values)
         elif com.is_integer_dtype(values.dtype):
             values = com._ensure_float64(values)
-        elif com.is_timedelta64_dtype(values.dtype):
-            values = com._ensure_float64(values.view('i8'))
+        elif com.needs_i8_conversion(values.dtype):
+            raise NotImplementedError("ops for {action} for this "
+                                      "dtype {dtype} are not "
+                                      "implemented".format(
+                                          action=self._window_type,
+                                          dtype=values.dtype))
         else:
             try:
                 values = com._ensure_float64(values)
@@ -498,15 +506,25 @@ def count(self):
         window = self._get_window()
         window = min(window, len(obj)) if not self.center else window
 
-        try:
-            converted = np.isfinite(obj).astype(float)
-        except TypeError:
-            converted = np.isfinite(obj.astype(float)).astype(float)
-        result = self._constructor(converted, window=window, min_periods=0,
-                                   center=self.center).sum()
+        blocks, obj = self._create_blocks(how=None)
+        results = []
+        for b in blocks:
 
-        result[result.isnull()] = 0
-        return result
+            if com.needs_i8_conversion(b.values):
+                result = b.notnull().astype(int)
+            else:
+                try:
+                    result = np.isfinite(b).astype(float)
+                except TypeError:
+                    result = np.isfinite(b.astype(float)).astype(float)
+
+                result[pd.isnull(result)] = 0
+
+            result = self._constructor(result, window=window, min_periods=0,
+                                       center=self.center).sum()
+            results.append(result)
+
+        return self._wrap_results(results, blocks, obj)
 
     _shared_docs['apply'] = dedent("""
     %(name)s function apply

pandas/tests/test_window.py

@@ -3,6 +3,7 @@
 import sys
 import warnings
 
+from nose.tools import assert_raises
 from datetime import datetime
 from numpy.random import randn
 from numpy.testing.decorators import slow
@@ -98,19 +99,6 @@ def tests_skip_nuisance(self):
         result = r.sum()
         assert_frame_equal(result, expected)
 
-    def test_timedeltas(self):
-
-        df = DataFrame({'A': range(5),
-                        'B': pd.timedelta_range('1 day', periods=5)})
-        r = df.rolling(window=3)
-        result = r.sum()
-        expected = DataFrame({'A': [np.nan, np.nan, 3, 6, 9],
-                              'B': pd.to_timedelta([pd.NaT, pd.NaT,
-                                                    '6 days', '9 days',
-                                                    '12 days'])},
-                             columns=list('AB'))
-        assert_frame_equal(result, expected)
-
     def test_agg(self):
         df = DataFrame({'A': range(5), 'B': range(0, 10, 2)})
 
@@ -291,8 +279,13 @@ def test_deprecations(self):
 
 # GH #12373 : rolling functions error on float32 data
 # make sure rolling functions works for different dtypes
-class TestDtype(Base):
-    dtype = None
+#
+# NOTE that these are yielded tests and so _create_data is
+# explicity called, nor do these inherit from unittest.TestCase
+#
+# further note that we are only checking rolling for fully dtype
+# compliance (though both expanding and ewm inherit)
+class Dtype(object):
     window = 2
 
     funcs = {
@@ -371,76 +364,84 @@ def _create_dtype_data(self, dtype):
         return data
 
     def _create_data(self):
-        super(TestDtype, self)._create_data()
         self.data = self._create_dtype_data(self.dtype)
         self.expects = self.get_expects()
 
-    def setUp(self):
-        self._create_data()
-
     def test_dtypes(self):
+        self._create_data()
         for f_name, d_name in product(self.funcs.keys(), self.data.keys()):
             f = self.funcs[f_name]
             d = self.data[d_name]
-            assert_equal = assert_series_equal if isinstance(
-                d, Series) else assert_frame_equal
             exp = self.expects[d_name][f_name]
+            yield self.check_dtypes, f, f_name, d, d_name, exp
 
-            roll = d.rolling(window=self.window)
-            result = f(roll)
-
-            assert_equal(result, exp)
+    def check_dtypes(self, f, f_name, d, d_name, exp):
+        roll = d.rolling(window=self.window)
+        result = f(roll)
+        assert_almost_equal(result, exp)
 
 
-class TestDtype_object(TestDtype):
+class TestDtype_object(Dtype):
     dtype = object
 
 
-class TestDtype_int8(TestDtype):
+class Dtype_integer(Dtype):
+    pass
+
+
+class TestDtype_int8(Dtype_integer):
     dtype = np.int8
 
 
-class TestDtype_int16(TestDtype):
+class TestDtype_int16(Dtype_integer):
     dtype = np.int16
 
 
-class TestDtype_int32(TestDtype):
+class TestDtype_int32(Dtype_integer):
     dtype = np.int32
 
 
-class TestDtype_int64(TestDtype):
+class TestDtype_int64(Dtype_integer):
     dtype = np.int64
 
 
-class TestDtype_uint8(TestDtype):
+class Dtype_uinteger(Dtype):
+    pass
+
+
+class TestDtype_uint8(Dtype_uinteger):
     dtype = np.uint8
 
 
-class TestDtype_uint16(TestDtype):
+class TestDtype_uint16(Dtype_uinteger):
     dtype = np.uint16
 
 
-class TestDtype_uint32(TestDtype):
+class TestDtype_uint32(Dtype_uinteger):
     dtype = np.uint32
 
 
-class TestDtype_uint64(TestDtype):
+class TestDtype_uint64(Dtype_uinteger):
     dtype = np.uint64
 
 
-class TestDtype_float16(TestDtype):
+class Dtype_float(Dtype):
+    pass
+
+
+class TestDtype_float16(Dtype_float):
     dtype = np.float16
 
 
-class TestDtype_float32(TestDtype):
+class TestDtype_float32(Dtype_float):
     dtype = np.float32
 
 
-class TestDtype_float64(TestDtype):
+class TestDtype_float64(Dtype_float):
     dtype = np.float64
 
 
-class TestDtype_category(TestDtype):
+class TestDtype_category(Dtype):
     dtype = 'category'
     include_df = False
 
@@ -456,25 +457,37 @@ def _create_dtype_data(self, dtype):
         return data
 
 
-class TestDatetimeLikeDtype(TestDtype):
-    dtype = np.dtype('M8[ns]')
+class DatetimeLike(Dtype):
 
-    # GH #12373: rolling functions raise ValueError on float32 data
-    def setUp(self):
-        raise nose.SkipTest("Skip rolling on DatetimeLike dtypes [{0}].".format(self.dtype))
+    def check_dtypes(self, f, f_name, d, d_name, exp):
 
-    def test_dtypes(self):
-        with tm.assertRaises(TypeError):
-            super(TestDatetimeLikeDtype, self).test_dtypes()
+        roll = d.rolling(window=self.window)
+
+        if f_name == 'count':
+            result = f(roll)
+            assert_almost_equal(result, exp)
 
+        else:
+
+            # other methods not Implemented ATM
+            assert_raises(NotImplementedError, f, roll)
 
-class TestDtype_timedelta(TestDatetimeLikeDtype):
+
+class TestDtype_timedelta(DatetimeLike):
     dtype = np.dtype('m8[ns]')
 
 
-class TestDtype_datetime64UTC(TestDatetimeLikeDtype):
+class TestDtype_datetime(DatetimeLike):
+    dtype = np.dtype('M8[ns]')
+
+
+class TestDtype_datetime64UTC(DatetimeLike):
     dtype = 'datetime64[ns, UTC]'
 
+    def _create_data(self):
+        raise nose.SkipTest("direct creation of extension dtype "
+                            "datetime64[ns, UTC] is not supported ATM")
+
 
 class TestMoments(Base):