Parametrize test for many datetimelikes

Author: watercrossing

pandas/tests/groupby/test_counting.py

@@ -3,7 +3,8 @@
 
 import numpy as np
 
-from pandas import (DataFrame, Series, MultiIndex, Timestamp)
+from pandas import (DataFrame, Series, MultiIndex, Timestamp, Timedelta,
+                    Period)
 from pandas.util.testing import (assert_series_equal, assert_frame_equal)
 from pandas.compat import (range, product as cart_product)
 
@@ -196,13 +197,27 @@ def test_ngroup_respects_groupby_order(self):
             assert_series_equal(Series(df['group_index'].values),
                                 g.cumcount())
 
-    def test_count_with_datetime(self):
-        # test for #13393
-        df = DataFrame({'x': ['a', 'a', 'b'],
-                        'y': [Timestamp('2016-05-07 20:09:25+00:00'),
-                              Timestamp('2016-05-07 20:09:29+00:00'),
-                              Timestamp('2016-05-07 20:09:29+00:00')]})
+    def test_count_with_datetimelike(self):
+        # test for #13393, where DataframeGroupBy.count() fails
+        # when counting a datetimelike column.
+
         expected = DataFrame({'y': [2, 1]}, index=['a', 'b'])
         expected.index.name = "x"
-        res = df.groupby('x').count()
-        assert_frame_equal(expected, res)
+
+        def _test(datetimelike):
+            df = DataFrame({'x': ['a', 'a', 'b'], 'y': datetimelike})
+            res = df.groupby('x').count()
+            assert_frame_equal(expected, res)
+
+        d1 = [Timestamp('2016-05-07 20:09:25+00:00'),
+              Timestamp('2016-05-07 20:09:29+00:00'),
+              Timestamp('2016-05-07 20:09:29+00:00')]
+        _test(d1)
+        d2 = [Timedelta(x, unit="h") for x in range(1, 4)]
+        _test(d2)
+        d3 = [Period(freq="2W", year=2017, month=x) for x in range(1, 4)]
+        _test(d3)
+        d4 = [Timestamp('2016-05-07 20:09:25'),
+              Timestamp('2016-05-07 20:09:29'),
+              Timestamp('2016-05-07 20:09:29')]
+        _test(d4)