TST: Extend timedelta64 arithmetic tests to TimedeltaArray

pandas/core/arrays/timedeltas.py

@@ -1,5 +1,8 @@
 # -*- coding: utf-8 -*-
+from __future__ import division
+
 from datetime import timedelta
+import operator
 import warnings
 
 import numpy as np
@@ -22,11 +25,13 @@
     is_datetime64_dtype,
     is_list_like,
     ensure_int64)
-from pandas.core.dtypes.generic import ABCSeries, ABCTimedeltaIndex
+from pandas.core.dtypes.generic import (
+    ABCDataFrame, ABCSeries, ABCTimedeltaIndex, ABCIndexClass)
 from pandas.core.dtypes.missing import isna
 
 import pandas.core.common as com
 from pandas.core.algorithms import checked_add_with_arr
+from pandas.core import ops
 
 from pandas.tseries.offsets import Tick
 from pandas.tseries.frequencies import to_offset
@@ -107,8 +112,32 @@ def wrapper(self, other):
     return compat.set_function_name(wrapper, opname, cls)
 
 
+def _wrap_tdi_op(op):
+    """
+    Instead of re-implementing multiplication/division etc operations
+    in the Array class, for now we dispatch to the TimedeltaIndex
+    implementations.
+    """
+    # TODO: implement directly here and wrap in TimedeltaIndex, instead of
+    #  the other way around
+    def method(self, other):
+        if isinstance(other, (ABCSeries, ABCDataFrame, ABCIndexClass)):
+            return NotImplemented
+
+        from pandas import TimedeltaIndex
+        obj = TimedeltaIndex(self)
+        result = op(obj, other)
+        if is_timedelta64_dtype(result):
+            return type(self)(result)
+        return np.array(result)
+
+    method.__name__ = '__{name}__'.format(name=op.__name__)
+    return method
+
+
 class TimedeltaArrayMixin(dtl.DatetimeLikeArrayMixin):
     _typ = "timedeltaarray"
+    __array_priority__ = 1000
 
     @property
     def _box_func(self):
@@ -288,6 +317,21 @@ def _evaluate_with_timedelta_like(self, other, op):
 
         return NotImplemented
 
+    __mul__ = _wrap_tdi_op(operator.mul)
+    __rmul__ = __mul__
+    __truediv__ = _wrap_tdi_op(operator.truediv)
+    __floordiv__ = _wrap_tdi_op(operator.floordiv)
+    __rfloordiv__ = _wrap_tdi_op(ops.rfloordiv)
+
+    if compat.PY2:
+        __div__ = __truediv__
+
+    # Note: TimedeltaIndex overrides this in call to cls._add_numeric_methods
+    def __neg__(self):
+        if self.freq is not None:
+            return type(self)(-self._data, freq=-self.freq)
+        return type(self)(-self._data)
+
     # ----------------------------------------------------------------
     # Conversion Methods - Vectorized analogues of Timedelta methods
 

pandas/core/indexes/base.py

@@ -19,6 +19,7 @@
     ABCSeries, ABCDataFrame,
     ABCMultiIndex,
     ABCPeriodIndex, ABCTimedeltaIndex, ABCDatetimeIndex,
+    ABCTimedeltaArray,
     ABCDateOffset)
 from pandas.core.dtypes.missing import isna, array_equivalent
 from pandas.core.dtypes.cast import maybe_cast_to_integer_array
@@ -123,7 +124,8 @@ def index_arithmetic_method(self, other):
         elif isinstance(other, ABCTimedeltaIndex):
             # Defer to subclass implementation
             return NotImplemented
-        elif isinstance(other, np.ndarray) and is_timedelta64_dtype(other):
+        elif (isinstance(other, (np.ndarray, ABCTimedeltaArray)) and
+              is_timedelta64_dtype(other)):
             # GH#22390; wrap in Series for op, this will in turn wrap in
             # TimedeltaIndex, but will correctly raise TypeError instead of
             # NullFrequencyError for add/sub ops

pandas/core/indexes/timedeltas.py

@@ -227,6 +227,14 @@ def _format_native_types(self, na_rep=u'NaT', date_format=None, **kwargs):
     # -------------------------------------------------------------------
     # Wrapping TimedeltaArray
 
+    __mul__ = Index.__mul__
+    __rmul__ = Index.__rmul__
+    __truediv__ = Index.__truediv__
+    __floordiv__ = Index.__floordiv__
+    __rfloordiv__ = Index.__rfloordiv__
+    if compat.PY2:
+        __div__ = Index.__div__
+
     days = wrap_field_accessor(TimedeltaArray.days)
     seconds = wrap_field_accessor(TimedeltaArray.seconds)
     microseconds = wrap_field_accessor(TimedeltaArray.microseconds)

pandas/tests/arithmetic/conftest.py

@@ -5,7 +5,10 @@
 import pandas as pd
 
 from pandas.compat import long
-from pandas.core.arrays import PeriodArray, DatetimeArrayMixin as DatetimeArray
+from pandas.core.arrays import (
+    PeriodArray,
+    DatetimeArrayMixin as DatetimeArray,
+    TimedeltaArrayMixin as TimedeltaArray)
 
 
 @pytest.fixture(params=[1, np.array(1, dtype=np.int64)])
@@ -189,3 +192,12 @@ def box_with_datetime(request):
     Like `box`, but specific to datetime64 for also testing DatetimeArray
     """
     return request.param
+
+
+@pytest.fixture(params=[pd.Index, pd.Series, pd.DataFrame, TimedeltaArray],
+                ids=lambda x: x.__name__)
+def box_with_timedelta(request):
+    """
+    Like `box`, but specific to timedelta64 for also testing TimedeltaArray
+    """
+    return request.param