Add CFTimeIndex.shift

doc/whats-new.rst

@@ -36,9 +36,15 @@ Enhancements
 - Added support for Python 3.7. (:issue:`2271`).
   By `Joe Hamman <https://github.com/jhamman>`_.
 
+- Added :py:meth:`~xarray.CFTimeIndex.shift` for shifting the values of a
+  CFTimeIndex by a specified frequency.  (:issue:`2244`).  By `Spencer Clark
+  <https://github.com/spencerkclark>`_.
+
 Bug fixes
 ~~~~~~~~~
 
+- Addition and subtraction operators used with a CFTimeIndex now preserve the
+  index's type. (:issue:`2244`). By `Spencer Clark <https://github.com/spencerkclark>`_.
 
 .. _whats-new.0.10.9:
 

xarray/coding/cftimeindex.py

@@ -314,3 +314,55 @@ def __contains__(self, key):
     def contains(self, key):
         """Needed for .loc based partial-string indexing"""
         return self.__contains__(key)
+
+    def shift(self, n, freq):
+        """Shift the CFTimeIndex a multiple of the given frequency.
+
+        See the documentation for :py:func:`~xarray.cftime_range()` for a
+        complete listing of valid frequency strings.
+
+        Parameters
+        ----------
+        n : int
+            Periods to shift by
+        freq : str or datetime.timedelta
+            A frequency string or datetime.timedelta object to shift by
+
+        Returns
+        -------
+        CFTimeIndex
+
+        See also
+        --------
+        pandas.DatetimeIndex.shift
+
+        Examples
+        --------
+
+        >>> index = xr.cftime_range('2000', periods=1, freq='M')
+        >>> index
+        CFTimeIndex([2000-01-31 00:00:00], dtype='object')
+        >>> index.shift(1, 'M')
+        CFTimeIndex([2000-02-29 00:00:00], dtype='object')
+        """
+        from .cftime_offsets import to_offset
+
+        if not isinstance(n, int):
+            raise TypeError("'n' must be an int, got {}.".format(n))
+        if isinstance(freq, timedelta):
+            return self + n * freq
+        elif isinstance(freq, pycompat.basestring):
+            return self + n * to_offset(freq)
+        else:
+            raise TypeError(
+                "'freq' must be of type "
+                "str or datetime.timedelta, got {}.".format(freq))
+
+    def __add__(self, other):
+        return CFTimeIndex(np.array(self) + other)
+
+    def __radd__(self, other):
+        return CFTimeIndex(other + np.array(self))
+
+    def __sub__(self, other):
+        return CFTimeIndex(np.array(self) - other)

xarray/tests/test_cftimeindex.py

@@ -616,3 +616,69 @@ def test_concat_cftimeindex(date_type, enable_cftimeindex):
 def test_empty_cftimeindex():
     index = CFTimeIndex([])
     assert index.date_type is None
+
+
+@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
+def test_cftimeindex_add(index):
+    date_type = index.date_type
+    expected_dates = [date_type(1, 1, 2), date_type(1, 2, 2),
+                      date_type(2, 1, 2), date_type(2, 2, 2)]
+    expected = CFTimeIndex(expected_dates)
+    result = index + timedelta(days=1)
+    assert result.equals(expected)
+    assert isinstance(result, CFTimeIndex)
+
+
+@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
+def test_cftimeindex_radd(index):
+    date_type = index.date_type
+    expected_dates = [date_type(1, 1, 2), date_type(1, 2, 2),
+                      date_type(2, 1, 2), date_type(2, 2, 2)]
+    expected = CFTimeIndex(expected_dates)
+    result = timedelta(days=1) + index
+    assert result.equals(expected)
+    assert isinstance(result, CFTimeIndex)
+
+
+@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
+def test_cftimeindex_sub(index):
+    date_type = index.date_type
+    expected_dates = [date_type(1, 1, 2), date_type(1, 2, 2),
+                      date_type(2, 1, 2), date_type(2, 2, 2)]
+    expected = CFTimeIndex(expected_dates)
+    result = index + timedelta(days=2)
+    result = result - timedelta(days=1)
+    assert result.equals(expected)
+    assert isinstance(result, CFTimeIndex)
+
+
+@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
+def test_cftimeindex_rsub(index):
+    with pytest.raises(TypeError):
+        timedelta(days=1) - index
+
+
+@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
+@pytest.mark.parametrize('freq', ['D', timedelta(days=1)])
+def test_cftimeindex_shift(index, freq):
+    date_type = index.date_type
+    expected_dates = [date_type(1, 1, 3), date_type(1, 2, 3),
+                      date_type(2, 1, 3), date_type(2, 2, 3)]
+    expected = CFTimeIndex(expected_dates)
+    result = index.shift(2, freq)
+    assert result.equals(expected)
+    assert isinstance(result, CFTimeIndex)
+
+
+@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
+def test_cftimeindex_shift_invalid_n():
+    index = xr.cftime_range('2000', periods=3)
+    with pytest.raises(TypeError):
+        index.shift('a', 'D')
+
+
+@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
+def test_cftimeindex_shift_invalid_freq():
+    index = xr.cftime_range('2000', periods=3)
+    with pytest.raises(TypeError):
+        index.shift(1, 1)