TST: Extend timedelta64 arithmetic tests to TimedeltaArray (pandas-dev#23642)

Author: jbrockmendel

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
@@ -16,9 +19,11 @@
     _TD_DTYPE, ensure_int64, is_datetime64_dtype, is_float_dtype,
     is_integer_dtype, is_list_like, is_object_dtype, is_string_dtype,
     is_timedelta64_dtype)
-from pandas.core.dtypes.generic import ABCSeries, ABCTimedeltaIndex
+from pandas.core.dtypes.generic import (
+    ABCDataFrame, ABCIndexClass, ABCSeries, ABCTimedeltaIndex)
 from pandas.core.dtypes.missing import isna
 
+from pandas.core import ops
 from pandas.core.algorithms import checked_add_with_arr
 import pandas.core.common as com
 
@@ -101,8 +106,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):
@@ -293,11 +322,25 @@ 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)
 
+    def __abs__(self):
+        # Note: freq is not preserved
+        return type(self)(np.abs(self._data))
+
     # ----------------------------------------------------------------
     # Conversion Methods - Vectorized analogues of Timedelta methods
 

pandas/core/indexes/base.py

@@ -19,7 +19,7 @@
     ABCSeries, ABCDataFrame,
     ABCMultiIndex,
     ABCPeriodIndex, ABCTimedeltaIndex, ABCDatetimeIndex,
-    ABCDateOffset, ABCIndexClass)
+    ABCDateOffset, ABCIndexClass, ABCTimedeltaArray)
 from pandas.core.dtypes.missing import isna, array_equivalent
 from pandas.core.dtypes.cast import maybe_cast_to_integer_array
 from pandas.core.dtypes.common import (
@@ -123,7 +123,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)