|
2 | 2 | from __future__ import print_function
|
3 | 3 |
|
4 | 4 | import numpy as np
|
| 5 | +import pytest |
5 | 6 |
|
6 |
| -from pandas import (DataFrame, Series, MultiIndex) |
7 |
| -from pandas.util.testing import assert_series_equal |
| 7 | +from pandas import (DataFrame, Series, MultiIndex, Timestamp, Timedelta, |
| 8 | + Period) |
| 9 | +from pandas.util.testing import (assert_series_equal, assert_frame_equal) |
8 | 10 | from pandas.compat import (range, product as cart_product)
|
9 | 11 |
|
10 | 12 |
|
@@ -195,3 +197,18 @@ def test_ngroup_respects_groupby_order(self):
|
195 | 197 | g.ngroup())
|
196 | 198 | assert_series_equal(Series(df['group_index'].values),
|
197 | 199 | g.cumcount())
|
| 200 | + |
| 201 | + @pytest.mark.parametrize('datetimelike', [ |
| 202 | + [Timestamp('2016-05-%02d 20:09:25+00:00' % i) for i in range(1, 4)], |
| 203 | + [Timestamp('2016-05-%02d 20:09:25' % i) for i in range(1, 4)], |
| 204 | + [Timedelta(x, unit="h") for x in range(1, 4)], |
| 205 | + [Period(freq="2W", year=2017, month=x) for x in range(1, 4)]]) |
| 206 | + def test_count_with_datetimelike(self, datetimelike): |
| 207 | + # test for #13393, where DataframeGroupBy.count() fails |
| 208 | + # when counting a datetimelike column. |
| 209 | + |
| 210 | + df = DataFrame({'x': ['a', 'a', 'b'], 'y': datetimelike}) |
| 211 | + res = df.groupby('x').count() |
| 212 | + expected = DataFrame({'y': [2, 1]}, index=['a', 'b']) |
| 213 | + expected.index.name = "x" |
| 214 | + assert_frame_equal(expected, res) |
0 commit comments