Skip to content

Commit 7cab33a

Browse files
spencerkclarkshoyer
authored andcommitted
Improve arithmetic operations involving CFTimeIndexes and TimedeltaIndexes (#2485)
* Improve arithmetic involving CFTimeIndexes and TimedeltaIndexes * Fix Appveyor Python 2.7 test failure * Clarify comment in _add_delta
1 parent 4bad455 commit 7cab33a

File tree

3 files changed

+67
-2
lines changed

3 files changed

+67
-2
lines changed

doc/whats-new.rst

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ Enhancements
7171
:py:meth:`~xarray.Dataset.differentiate`,
7272
:py:meth:`~xarray.DataArray.interp`, and
7373
:py:meth:`~xarray.Dataset.interp`.
74-
By `Spencer Clark <https://github.com/spencerkclark>`_
74+
By `Spencer Clark <https://github.com/spencerkclark>`_.
7575

7676
Bug fixes
7777
~~~~~~~~~
@@ -97,6 +97,13 @@ Bug fixes
9797
``open_rasterio`` to error (:issue:`2454`).
9898
By `Stephan Hoyer <https://github.com/shoyer>`_.
9999

100+
- Subtracting one CFTimeIndex from another now returns a
101+
``pandas.TimedeltaIndex``, analogous to the behavior for DatetimeIndexes
102+
(:issue:`2484`). By `Spencer Clark <https://github.com/spencerkclark>`_.
103+
- Adding a TimedeltaIndex to, or subtracting a TimedeltaIndex from a
104+
CFTimeIndex is now allowed (:issue:`2484`).
105+
By `Spencer Clark <https://github.com/spencerkclark>`_.
106+
100107
.. _whats-new.0.10.9:
101108

102109
v0.10.9 (21 September 2018)

xarray/coding/cftimeindex.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -359,13 +359,27 @@ def shift(self, n, freq):
359359
"str or datetime.timedelta, got {}.".format(freq))
360360

361361
def __add__(self, other):
362+
if isinstance(other, pd.TimedeltaIndex):
363+
other = other.to_pytimedelta()
362364
return CFTimeIndex(np.array(self) + other)
363365

364366
def __radd__(self, other):
367+
if isinstance(other, pd.TimedeltaIndex):
368+
other = other.to_pytimedelta()
365369
return CFTimeIndex(other + np.array(self))
366370

367371
def __sub__(self, other):
368-
return CFTimeIndex(np.array(self) - other)
372+
if isinstance(other, CFTimeIndex):
373+
return pd.TimedeltaIndex(np.array(self) - np.array(other))
374+
elif isinstance(other, pd.TimedeltaIndex):
375+
return CFTimeIndex(np.array(self) - other.to_pytimedelta())
376+
else:
377+
return CFTimeIndex(np.array(self) - other)
378+
379+
def _add_delta(self, deltas):
380+
# To support TimedeltaIndex + CFTimeIndex with older versions of
381+
# pandas. No longer used as of pandas 0.23.
382+
return self + deltas
369383

370384

371385
def _parse_iso8601_without_reso(date_type, datetime_str):

xarray/tests/test_cftimeindex.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -629,6 +629,17 @@ def test_cftimeindex_add(index):
629629
assert isinstance(result, CFTimeIndex)
630630

631631

632+
@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
633+
@pytest.mark.parametrize('calendar', _CFTIME_CALENDARS)
634+
def test_cftimeindex_add_timedeltaindex(calendar):
635+
a = xr.cftime_range('2000', periods=5, calendar=calendar)
636+
deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)])
637+
result = a + deltas
638+
expected = a.shift(2, 'D')
639+
assert result.equals(expected)
640+
assert isinstance(result, CFTimeIndex)
641+
642+
632643
@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
633644
def test_cftimeindex_radd(index):
634645
date_type = index.date_type
@@ -640,6 +651,17 @@ def test_cftimeindex_radd(index):
640651
assert isinstance(result, CFTimeIndex)
641652

642653

654+
@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
655+
@pytest.mark.parametrize('calendar', _CFTIME_CALENDARS)
656+
def test_timedeltaindex_add_cftimeindex(calendar):
657+
a = xr.cftime_range('2000', periods=5, calendar=calendar)
658+
deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)])
659+
result = deltas + a
660+
expected = a.shift(2, 'D')
661+
assert result.equals(expected)
662+
assert isinstance(result, CFTimeIndex)
663+
664+
643665
@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
644666
def test_cftimeindex_sub(index):
645667
date_type = index.date_type
@@ -652,6 +674,28 @@ def test_cftimeindex_sub(index):
652674
assert isinstance(result, CFTimeIndex)
653675

654676

677+
@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
678+
@pytest.mark.parametrize('calendar', _CFTIME_CALENDARS)
679+
def test_cftimeindex_sub_cftimeindex(calendar):
680+
a = xr.cftime_range('2000', periods=5, calendar=calendar)
681+
b = a.shift(2, 'D')
682+
result = b - a
683+
expected = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)])
684+
assert result.equals(expected)
685+
assert isinstance(result, pd.TimedeltaIndex)
686+
687+
688+
@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
689+
@pytest.mark.parametrize('calendar', _CFTIME_CALENDARS)
690+
def test_cftimeindex_sub_timedeltaindex(calendar):
691+
a = xr.cftime_range('2000', periods=5, calendar=calendar)
692+
deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)])
693+
result = a - deltas
694+
expected = a.shift(-2, 'D')
695+
assert result.equals(expected)
696+
assert isinstance(result, CFTimeIndex)
697+
698+
655699
@pytest.mark.skipif(not has_cftime, reason='cftime not installed')
656700
def test_cftimeindex_rsub(index):
657701
with pytest.raises(TypeError):

0 commit comments

Comments
 (0)