diff --git a/doc/api-hidden.rst b/doc/api-hidden.rst index 1826cc86892..0e8143c72ea 100644 --- a/doc/api-hidden.rst +++ b/doc/api-hidden.rst @@ -151,3 +151,5 @@ plot.FacetGrid.set_titles plot.FacetGrid.set_ticks plot.FacetGrid.map + + CFTimeIndex.shift diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 8b145924f2d..3f8d40910cd 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -46,6 +46,9 @@ Enhancements - Added support for Python 3.7. (:issue:`2271`). By `Joe Hamman `_. +- Added :py:meth:`~xarray.CFTimeIndex.shift` for shifting the values of a + CFTimeIndex by a specified frequency. (:issue:`2244`). By `Spencer Clark + `_. - 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 `_. - ``xarray.DataArray.roll`` correctly handles multidimensional arrays. (:issue:`2445`) By `Keisuke Fujii `_. diff --git a/xarray/coding/cftimeindex.py b/xarray/coding/cftimeindex.py index faf1a044505..341ecfed262 100644 --- a/xarray/coding/cftimeindex.py +++ b/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) diff --git a/xarray/tests/test_cftimeindex.py b/xarray/tests/test_cftimeindex.py index 62a29a15247..a558ab9a784 100644 --- a/xarray/tests/test_cftimeindex.py +++ b/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