POC: _eadata (#24394)

Author: jbrockmendel

Diff for: pandas/core/indexes/datetimelike.py

@@ -39,7 +39,6 @@ class DatetimeIndexOpsMixin(DatetimeLikeArrayMixin):
 
     # override DatetimeLikeArrayMixin method
     copy = Index.copy
-    unique = Index.unique
 
     # DatetimeLikeArrayMixin assumes subclasses are mutable, so these are
     # properties there.  They can be made into cache_readonly for Index
@@ -51,6 +50,30 @@ class DatetimeIndexOpsMixin(DatetimeLikeArrayMixin):
     _resolution = cache_readonly(DatetimeLikeArrayMixin._resolution.fget)
     resolution = cache_readonly(DatetimeLikeArrayMixin.resolution.fget)
 
+    def unique(self, level=None):
+        if level is not None:
+            self._validate_index_level(level)
+
+        result = self._eadata.unique()
+
+        # Note: if `self` is already unique, then self.unique() should share
+        #  a `freq` with self.  If not already unique, then self.freq must be
+        #  None, so again sharing freq is correct.
+        return self._shallow_copy(result._data)
+
+    @classmethod
+    def _create_comparison_method(cls, op):
+        """
+        Create a comparison method that dispatches to ``cls.values``.
+        """
+        def wrapper(self, other):
+            result = op(self._eadata, maybe_unwrap_index(other))
+            return result
+
+        wrapper.__doc__ = op.__doc__
+        wrapper.__name__ = '__{}__'.format(op.__name__)
+        return wrapper
+
     # A few methods that are shared
     _maybe_mask_results = DatetimeLikeArrayMixin._maybe_mask_results
 
@@ -106,7 +129,7 @@ def wrapper(left, right):
 
     @Appender(DatetimeLikeArrayMixin._evaluate_compare.__doc__)
     def _evaluate_compare(self, other, op):
-        result = DatetimeLikeArrayMixin._evaluate_compare(self, other, op)
+        result = self._eadata._evaluate_compare(other, op)
         if is_bool_dtype(result):
             return result
         try:
@@ -406,7 +429,7 @@ def _add_datetimelike_methods(cls):
 
         def __add__(self, other):
             # dispatch to ExtensionArray implementation
-            result = super(cls, self).__add__(other)
+            result = self._eadata.__add__(maybe_unwrap_index(other))
             return wrap_arithmetic_op(self, other, result)
 
         cls.__add__ = __add__
@@ -418,13 +441,13 @@ def __radd__(self, other):
 
         def __sub__(self, other):
             # dispatch to ExtensionArray implementation
-            result = super(cls, self).__sub__(other)
+            result = self._eadata.__sub__(maybe_unwrap_index(other))
             return wrap_arithmetic_op(self, other, result)
 
         cls.__sub__ = __sub__
 
         def __rsub__(self, other):
-            result = super(cls, self).__rsub__(other)
+            result = self._eadata.__rsub__(maybe_unwrap_index(other))
             return wrap_arithmetic_op(self, other, result)
 
         cls.__rsub__ = __rsub__
@@ -548,9 +571,8 @@ def astype(self, dtype, copy=True):
 
     @Appender(DatetimeLikeArrayMixin._time_shift.__doc__)
     def _time_shift(self, periods, freq=None):
-        result = DatetimeLikeArrayMixin._time_shift(self, periods, freq=freq)
-        result.name = self.name
-        return result
+        result = self._eadata._time_shift(periods, freq=freq)
+        return type(self)(result, name=self.name)
 
 
 def wrap_arithmetic_op(self, other, result):
@@ -589,7 +611,7 @@ def wrap_array_method(method, pin_name=False):
     method
     """
     def index_method(self, *args, **kwargs):
-        result = method(self, *args, **kwargs)
+        result = method(self._eadata, *args, **kwargs)
 
         # Index.__new__ will choose the appropriate subclass to return
         result = Index(result)
@@ -618,7 +640,7 @@ def wrap_field_accessor(prop):
     fget = prop.fget
 
     def f(self):
-        result = fget(self)
+        result = fget(self._eadata)
         if is_bool_dtype(result):
             # return numpy array b/c there is no BoolIndex
             return result
@@ -629,6 +651,28 @@ def f(self):
     return property(f)
 
 
+def maybe_unwrap_index(obj):
+    """
+    If operating against another Index object, we need to unwrap the underlying
+    data before deferring to the DatetimeArray/TimedeltaArray/PeriodArray
+    implementation, otherwise we will incorrectly return NotImplemented.
+
+    Parameters
+    ----------
+    obj : object
+
+    Returns
+    -------
+    unwrapped object
+    """
+    if isinstance(obj, ABCIndexClass):
+        if isinstance(obj, DatetimeIndexOpsMixin):
+            # i.e. PeriodIndex/DatetimeIndex/TimedeltaIndex
+            return obj._eadata
+        return obj._data
+    return obj
+
+
 class DatetimelikeDelegateMixin(PandasDelegate):
     """
     Delegation mechanism, specific for Datetime, Timedelta, and Period types.

Diff for: pandas/core/indexes/datetimes.py

@@ -713,15 +713,6 @@ def snap(self, freq='S'):
         return DatetimeIndex._simple_new(snapped, freq=freq)
         # TODO: what about self.name?  tz? if so, use shallow_copy?
 
-    def unique(self, level=None):
-        if level is not None:
-            self._validate_index_level(level)
-
-        # TODO(DatetimeArray): change dispatch once inheritance is removed
-        # call DatetimeArray method
-        result = DatetimeArray.unique(self)
-        return self._shallow_copy(result._data)
-
     def join(self, other, how='left', level=None, return_indexers=False,
              sort=False):
         """
@@ -1089,6 +1080,11 @@ def slice_indexer(self, start=None, end=None, step=None, kind=None):
     # --------------------------------------------------------------------
     # Wrapping DatetimeArray
 
+    @property
+    def _eadata(self):
+        return DatetimeArray._simple_new(self._data,
+                                         tz=self.tz, freq=self.freq)
+
     # Compat for frequency inference, see GH#23789
     _is_monotonic_increasing = Index.is_monotonic_increasing
     _is_monotonic_decreasing = Index.is_monotonic_decreasing

Diff for: pandas/core/indexes/period.py

@@ -26,7 +26,7 @@
 import pandas.core.indexes.base as ibase
 from pandas.core.indexes.base import _index_shared_docs, ensure_index
 from pandas.core.indexes.datetimelike import (
-    DatetimeIndexOpsMixin, DatetimelikeDelegateMixin, wrap_arithmetic_op)
+    DatetimeIndexOpsMixin, DatetimelikeDelegateMixin)
 from pandas.core.indexes.datetimes import DatetimeIndex, Index, Int64Index
 from pandas.core.missing import isna
 from pandas.core.ops import get_op_result_name
@@ -247,6 +247,10 @@ def _simple_new(cls, values, name=None, freq=None, **kwargs):
     # ------------------------------------------------------------------------
     # Data
 
+    @property
+    def _eadata(self):
+        return self._data
+
     @property
     def _ndarray_values(self):
         return self._data._ndarray_values
@@ -878,52 +882,6 @@ def __setstate__(self, state):
 
     _unpickle_compat = __setstate__
 
-    @classmethod
-    def _add_datetimelike_methods(cls):
-        """
-        add in the datetimelike methods (as we may have to override the
-        superclass)
-        """
-        # TODO(DatetimeArray): move this up to DatetimeArrayMixin
-
-        def __add__(self, other):
-            # dispatch to ExtensionArray implementation
-            result = self._data.__add__(other)
-            return wrap_arithmetic_op(self, other, result)
-
-        cls.__add__ = __add__
-
-        def __radd__(self, other):
-            # alias for __add__
-            return self.__add__(other)
-        cls.__radd__ = __radd__
-
-        def __sub__(self, other):
-            # dispatch to ExtensionArray implementation
-            result = self._data.__sub__(other)
-            return wrap_arithmetic_op(self, other, result)
-
-        cls.__sub__ = __sub__
-
-        def __rsub__(self, other):
-            result = self._data.__rsub__(other)
-            return wrap_arithmetic_op(self, other, result)
-
-        cls.__rsub__ = __rsub__
-
-    @classmethod
-    def _create_comparison_method(cls, op):
-        """
-        Create a comparison method that dispatches to ``cls.values``.
-        """
-        # TODO(DatetimeArray): move to base class.
-        def wrapper(self, other):
-            return op(self._data, other)
-
-        wrapper.__doc__ = op.__doc__
-        wrapper.__name__ = '__{}__'.format(op.__name__)
-        return wrapper
-
     def view(self, dtype=None, type=None):
         # TODO(DatetimeArray): remove
         if dtype is None or dtype is __builtins__['type'](self):

Diff for: pandas/core/indexes/timedeltas.py

@@ -22,8 +22,8 @@
 import pandas.core.common as com
 from pandas.core.indexes.base import Index, _index_shared_docs
 from pandas.core.indexes.datetimelike import (
-    DatetimeIndexOpsMixin, wrap_arithmetic_op, wrap_array_method,
-    wrap_field_accessor)
+    DatetimeIndexOpsMixin, maybe_unwrap_index, wrap_arithmetic_op,
+    wrap_array_method, wrap_field_accessor)
 from pandas.core.indexes.numeric import Int64Index
 from pandas.core.ops import get_op_result_name
 from pandas.core.tools.timedeltas import _coerce_scalar_to_timedelta_type
@@ -36,11 +36,7 @@ def _make_wrapped_arith_op(opname):
     meth = getattr(TimedeltaArray, opname)
 
     def method(self, other):
-        oth = other
-        if isinstance(other, Index):
-            oth = other._data
-
-        result = meth(self, oth)
+        result = meth(self._eadata, maybe_unwrap_index(other))
         return wrap_arithmetic_op(self, other, result)
 
     method.__name__ = opname
@@ -237,6 +233,10 @@ def _format_native_types(self, na_rep=u'NaT', date_format=None, **kwargs):
     # -------------------------------------------------------------------
     # Wrapping TimedeltaArray
 
+    @property
+    def _eadata(self):
+        return TimedeltaArray._simple_new(self._data, freq=self.freq)
+
     __mul__ = _make_wrapped_arith_op("__mul__")
     __rmul__ = _make_wrapped_arith_op("__rmul__")
     __floordiv__ = _make_wrapped_arith_op("__floordiv__")
@@ -245,6 +245,11 @@ def _format_native_types(self, na_rep=u'NaT', date_format=None, **kwargs):
     __rmod__ = _make_wrapped_arith_op("__rmod__")
     __divmod__ = _make_wrapped_arith_op("__divmod__")
     __rdivmod__ = _make_wrapped_arith_op("__rdivmod__")
+    __truediv__ = _make_wrapped_arith_op("__truediv__")
+    __rtruediv__ = _make_wrapped_arith_op("__rtruediv__")
+    if compat.PY2:
+        __div__ = __truediv__
+        __rdiv__ = __rtruediv__
 
     days = wrap_field_accessor(TimedeltaArray.days)
     seconds = wrap_field_accessor(TimedeltaArray.seconds)
@@ -253,26 +258,6 @@ def _format_native_types(self, na_rep=u'NaT', date_format=None, **kwargs):
 
     total_seconds = wrap_array_method(TimedeltaArray.total_seconds, True)
 
-    def __truediv__(self, other):
-        oth = other
-        if isinstance(other, Index):
-            # TimedeltaArray defers, so we need to unwrap
-            oth = other._values
-        result = TimedeltaArray.__truediv__(self, oth)
-        return wrap_arithmetic_op(self, other, result)
-
-    def __rtruediv__(self, other):
-        oth = other
-        if isinstance(other, Index):
-            # TimedeltaArray defers, so we need to unwrap
-            oth = other._values
-        result = TimedeltaArray.__rtruediv__(self, oth)
-        return wrap_arithmetic_op(self, other, result)
-
-    if compat.PY2:
-        __div__ = __truediv__
-        __rdiv__ = __rtruediv__
-
     # Compat for frequency inference, see GH#23789
     _is_monotonic_increasing = Index.is_monotonic_increasing
     _is_monotonic_decreasing = Index.is_monotonic_decreasing

Diff for: pandas/tests/arithmetic/test_datetime64.py

@@ -1591,7 +1591,8 @@ def check(get_ser, test_ser):
             # with 'operate' (from core/ops.py) for the ops that are not
             # defined
             op = getattr(get_ser, op_str, None)
-            with pytest.raises(TypeError, match='operate|[cC]annot'):
+            with pytest.raises(TypeError,
+                               match='operate|[cC]annot|unsupported operand'):
                 op(test_ser)
 
         # ## timedelta64 ###
@@ -1973,15 +1974,15 @@ def test_dti_sub_tdi(self, tz_naive_fixture):
         result = dti - tdi
         tm.assert_index_equal(result, expected)
 
-        msg = 'cannot subtract .*TimedeltaIndex'
+        msg = 'cannot subtract .*TimedeltaArrayMixin'
         with pytest.raises(TypeError, match=msg):
             tdi - dti
 
         # sub with timedelta64 array
         result = dti - tdi.values
         tm.assert_index_equal(result, expected)
 
-        msg = 'cannot subtract DatetimeIndex from'
+        msg = 'cannot subtract DatetimeArrayMixin from'
         with pytest.raises(TypeError, match=msg):
             tdi.values - dti
 
@@ -1997,7 +1998,7 @@ def test_dti_isub_tdi(self, tz_naive_fixture):
         result -= tdi
         tm.assert_index_equal(result, expected)
 
-        msg = 'cannot subtract .*TimedeltaIndex'
+        msg = 'cannot subtract .*TimedeltaArrayMixin'
         with pytest.raises(TypeError, match=msg):
             tdi -= dti
 
@@ -2008,7 +2009,7 @@ def test_dti_isub_tdi(self, tz_naive_fixture):
 
         msg = '|'.join(['cannot perform __neg__ with this index type:',
                         'ufunc subtract cannot use operands with types',
-                        'cannot subtract DatetimeIndex from'])
+                        'cannot subtract DatetimeArrayMixin from'])
         with pytest.raises(TypeError, match=msg):
             tdi.values -= dti
 
@@ -2028,7 +2029,9 @@ def test_dti_isub_tdi(self, tz_naive_fixture):
     def test_add_datetimelike_and_dti(self, addend, tz):
         # GH#9631
         dti = DatetimeIndex(['2011-01-01', '2011-01-02']).tz_localize(tz)
-        msg = 'cannot add DatetimeIndex and {0}'.format(type(addend).__name__)
+        msg = ('cannot add DatetimeArrayMixin and {0}'
+               .format(type(addend).__name__)).replace('DatetimeIndex',
+                                                       'DatetimeArrayMixin')
         with pytest.raises(TypeError, match=msg):
             dti + addend
         with pytest.raises(TypeError, match=msg):