Add set_freq to dt accessor

pandas-dev · May 3, 2018 · 5778e28 · 5778e28
1 parent b0db6c3
commit 5778e28
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 4 deletions.
diff --git a/pandas/core/indexes/datetimes.py b/pandas/core/indexes/datetimes.py
@@ -320,7 +320,7 @@ def _add_comparison_methods(cls):
     _other_ops = ['date', 'time']
     _datetimelike_ops = _field_ops + _object_ops + _bool_ops + _other_ops
     _datetimelike_methods = ['to_period', 'tz_localize',
-                             'tz_convert',
+                             'tz_convert', 'set_freq',
                              'normalize', 'strftime', 'round', 'floor',
                              'ceil', 'month_name', 'day_name']
 

diff --git a/pandas/core/indexes/timedeltas.py b/pandas/core/indexes/timedeltas.py
@@ -178,7 +178,7 @@ def _join_i8_wrapper(joinf, **kwargs):
     _field_ops = ['days', 'seconds', 'microseconds', 'nanoseconds']
     _datetimelike_ops = _field_ops + _object_ops + _bool_ops
     _datetimelike_methods = ["to_pytimedelta", "total_seconds",
-                             "round", "floor", "ceil"]
+                             "round", "floor", "ceil", "set_freq"]
 
     @classmethod
     def _add_comparison_methods(cls):

diff --git a/pandas/tests/series/test_api.py b/pandas/tests/series/test_api.py
@@ -705,6 +705,7 @@ def test_dt_accessor_api_for_categorical(self):
             ('floor', ("D",), {}),
             ('ceil', ("D",), {}),
             ('asfreq', ("D",), {}),
+            ('set_freq', ("24H",), {}),
             # ('tz_localize', ("UTC",), {}),
         ]
         _special_func_names = [f[0] for f in special_func_defs]

diff --git a/pandas/tests/series/test_datetime_values.py b/pandas/tests/series/test_datetime_values.py
@@ -34,10 +34,11 @@ def test_dt_namespace_accessor(self):
         ok_for_dt = DatetimeIndex._datetimelike_ops
         ok_for_dt_methods = ['to_period', 'to_pydatetime', 'tz_localize',
                              'tz_convert', 'normalize', 'strftime', 'round',
-                             'floor', 'ceil', 'day_name', 'month_name']
+                             'floor', 'ceil', 'day_name', 'month_name',
+                             'set_freq']
         ok_for_td = TimedeltaIndex._datetimelike_ops
         ok_for_td_methods = ['components', 'to_pytimedelta', 'total_seconds',
-                             'round', 'floor', 'ceil']
+                             'round', 'floor', 'ceil', 'set_freq']
 
         def get_expected(s, name):
             result = getattr(Index(s._values), prop)
@@ -460,3 +461,38 @@ def test_datetime_understood(self):
         expected = pd.Series(pd.to_datetime([
             '2011-12-26', '2011-12-27', '2011-12-28']))
         tm.assert_series_equal(result, expected)
+
+    @pytest.mark.xfail(reason='GH 20937')
+    @pytest.mark.parametrize('data, new_freq', [
+        (date_range('20180101', periods=3), 'B'),
+        (date_range('20180101', periods=3, tz='US/Eastern'), 'B'),
+        (timedelta_range('1 day', periods=3), '86400S')])
+    def test_set_freq(self, data, new_freq):
+        # GH 20886
+        s = Series(data)
+
+        # can set to alternate freq
+        result = s.dt.set_freq(new_freq)
+        assert result.dt.freq == new_freq
+
+        # can reset to None
+        result = s.dt.set_freq(None)
+        assert result.dt.freq is None
+
+    @pytest.mark.parametrize('data', [
+        date_range('20180101', periods=3),
+        date_range('20180101', periods=3, tz='US/Eastern'),
+        timedelta_range('1 day', periods=3)])
+    def test_set_freq_errors(self, data):
+        # GH 20886
+        s = Series(data)
+
+        # setting with an incompatible freq
+        msg = ('Inferred frequency D from passed values does not conform to '
+               'passed frequency 5D')
+        with tm.assert_raises_regex(ValueError, msg):
+            s.dt.set_freq('5D')
+
+        # setting with non-freq string
+        with tm.assert_raises_regex(ValueError, 'Invalid frequency'):
+            s.dt.set_freq('foo')