Add CFTimeIndex.shift (#2431)

* Add CFTimeIndex.shift

Update what's new

Add bug fix note

* Add example to docstring

* Use pycompat for basestring

* Generate an API reference page for CFTimeIndex.shift

doc/api-hidden.rst

@@ -151,3 +151,5 @@
    plot.FacetGrid.set_titles
    plot.FacetGrid.set_ticks
    plot.FacetGrid.map
+
+   CFTimeIndex.shift

doc/whats-new.rst

@@ -46,6 +46,9 @@ 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>`_.
 - Added support for using ``cftime.datetime`` coordinates with
   :py:meth:`~xarray.DataArray.differentiate`,
   :py:meth:`~xarray.Dataset.differentiate`,
@@ -56,6 +59,8 @@ Enhancements
 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>`_.
 - ``xarray.DataArray.roll`` correctly handles multidimensional arrays.
   (:issue:`2445`)
   By `Keisuke Fujii <https://github.com/fujiisoup>`_.

xarray/coding/cftimeindex.py

@@ -315,6 +315,57 @@ 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)
+
 
 def _parse_iso8601_without_reso(date_type, datetime_str):
     date, _ = _parse_iso8601_with_reso(date_type, datetime_str)

xarray/tests/test_cftimeindex.py

@@ -619,6 +619,72 @@ def test_empty_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)
+
+
 @requires_cftime
 def test_parse_array_of_cftime_strings():
     from cftime import DatetimeNoLeap