REF: move EA wrapping/unwrapping to indexes.extensions (#30648)

Author: jbrockmendel

pandas/core/indexes/category.py

@@ -29,6 +29,7 @@
 import pandas.core.common as com
 import pandas.core.indexes.base as ibase
 from pandas.core.indexes.base import Index, _index_shared_docs, maybe_extract_name
+from pandas.core.indexes.extension import make_wrapped_comparison_op
 import pandas.core.missing as missing
 from pandas.core.ops import get_op_result_name
 
@@ -876,14 +877,7 @@ def _add_comparison_methods(cls):
         def _make_compare(op):
             opname = f"__{op.__name__}__"
 
-            def _evaluate_compare(self, other):
-                with np.errstate(all="ignore"):
-                    result = op(self.array, other)
-                if isinstance(result, ABCSeries):
-                    # Dispatch to pd.Categorical returned NotImplemented
-                    # and we got a Series back; down-cast to ndarray
-                    result = result._values
-                return result
+            _evaluate_compare = make_wrapped_comparison_op(opname)
 
             return compat.set_function_name(_evaluate_compare, opname, cls)
 

pandas/core/indexes/datetimelike.py

@@ -25,7 +25,7 @@
 )
 from pandas.core.dtypes.generic import ABCIndex, ABCIndexClass, ABCSeries
 
-from pandas.core import algorithms, ops
+from pandas.core import algorithms
 from pandas.core.accessor import PandasDelegate
 from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin
 from pandas.core.arrays.datetimelike import (
@@ -40,21 +40,11 @@
 
 from pandas.tseries.frequencies import DateOffset, to_offset
 
-from .extension import inherit_names
+from .extension import inherit_names, make_wrapped_arith_op, make_wrapped_comparison_op
 
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
 
 
-def _make_wrapped_arith_op(opname):
-    def method(self, other):
-        meth = getattr(self._data, opname)
-        result = meth(maybe_unwrap_index(other))
-        return wrap_arithmetic_op(self, other, result)
-
-    method.__name__ = opname
-    return method
-
-
 def _join_i8_wrapper(joinf, with_indexers: bool = True):
     """
     Create the join wrapper methods.
@@ -125,19 +115,7 @@ def _create_comparison_method(cls, op):
         """
         Create a comparison method that dispatches to ``cls.values``.
         """
-
-        def wrapper(self, other):
-            if isinstance(other, ABCSeries):
-                # the arrays defer to Series for comparison ops but the indexes
-                #  don't, so we have to unwrap here.
-                other = other._values
-
-            result = op(self._data, maybe_unwrap_index(other))
-            return result
-
-        wrapper.__doc__ = op.__doc__
-        wrapper.__name__ = f"__{op.__name__}__"
-        return wrapper
+        return make_wrapped_comparison_op(f"__{op.__name__}__")
 
     # ------------------------------------------------------------------------
     # Abstract data attributes
@@ -467,22 +445,22 @@ def _convert_scalar_indexer(self, key, kind=None):
 
         return super()._convert_scalar_indexer(key, kind=kind)
 
-    __add__ = _make_wrapped_arith_op("__add__")
-    __radd__ = _make_wrapped_arith_op("__radd__")
-    __sub__ = _make_wrapped_arith_op("__sub__")
-    __rsub__ = _make_wrapped_arith_op("__rsub__")
-    __pow__ = _make_wrapped_arith_op("__pow__")
-    __rpow__ = _make_wrapped_arith_op("__rpow__")
-    __mul__ = _make_wrapped_arith_op("__mul__")
-    __rmul__ = _make_wrapped_arith_op("__rmul__")
-    __floordiv__ = _make_wrapped_arith_op("__floordiv__")
-    __rfloordiv__ = _make_wrapped_arith_op("__rfloordiv__")
-    __mod__ = _make_wrapped_arith_op("__mod__")
-    __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__")
+    __add__ = make_wrapped_arith_op("__add__")
+    __radd__ = make_wrapped_arith_op("__radd__")
+    __sub__ = make_wrapped_arith_op("__sub__")
+    __rsub__ = make_wrapped_arith_op("__rsub__")
+    __pow__ = make_wrapped_arith_op("__pow__")
+    __rpow__ = make_wrapped_arith_op("__rpow__")
+    __mul__ = make_wrapped_arith_op("__mul__")
+    __rmul__ = make_wrapped_arith_op("__rmul__")
+    __floordiv__ = make_wrapped_arith_op("__floordiv__")
+    __rfloordiv__ = make_wrapped_arith_op("__rfloordiv__")
+    __mod__ = make_wrapped_arith_op("__mod__")
+    __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__")
 
     def isin(self, values, level=None):
         """
@@ -864,55 +842,13 @@ def _wrap_joined_index(self, joined, other):
             return self._simple_new(joined, name, **kwargs)
 
 
-def wrap_arithmetic_op(self, other, result):
-    if result is NotImplemented:
-        return NotImplemented
-
-    if isinstance(result, tuple):
-        # divmod, rdivmod
-        assert len(result) == 2
-        return (
-            wrap_arithmetic_op(self, other, result[0]),
-            wrap_arithmetic_op(self, other, result[1]),
-        )
-
-    if not isinstance(result, Index):
-        # Index.__new__ will choose appropriate subclass for dtype
-        result = Index(result)
-
-    res_name = ops.get_op_result_name(self, other)
-    result.name = res_name
-    return result
-
-
-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):
-        return obj._data
-    return obj
-
-
 class DatetimelikeDelegateMixin(PandasDelegate):
     """
     Delegation mechanism, specific for Datetime, Timedelta, and Period types.
 
     Functionality is delegated from the Index class to an Array class. A
     few things can be customized
 
-    * _delegate_class : type
-        The class being delegated to.
     * _delegated_methods, delegated_properties : List
         The list of property / method names being delagated.
     * raw_methods : Set
@@ -929,10 +865,6 @@ class DatetimelikeDelegateMixin(PandasDelegate):
     _raw_properties: Set[str] = set()
     _data: ExtensionArray
 
-    @property
-    def _delegate_class(self):
-        raise AbstractMethodError
-
     def _delegate_property_get(self, name, *args, **kwargs):
         result = getattr(self._data, name)
         if name not in self._raw_properties:

pandas/core/indexes/datetimes.py

@@ -86,7 +86,6 @@ class DatetimeDelegateMixin(DatetimelikeDelegateMixin):
         | set(_extra_raw_properties)
     )
     _raw_methods = set(_extra_raw_methods)
-    _delegate_class = DatetimeArray
 
 
 @inherit_names(["_timezone", "is_normalized", "_resolution"], DatetimeArray, cache=True)

pandas/core/indexes/extension.py

@@ -5,6 +5,12 @@
 
 from pandas.util._decorators import cache_readonly
 
+from pandas.core.dtypes.generic import ABCSeries
+
+from pandas.core.ops import get_op_result_name
+
+from .base import Index
+
 
 def inherit_from_data(name: str, delegate, cache: bool = False):
     """
@@ -76,3 +82,73 @@ def wrapper(cls):
         return cls
 
     return wrapper
+
+
+def make_wrapped_comparison_op(opname):
+    """
+    Create a comparison method that dispatches to ``._data``.
+    """
+
+    def wrapper(self, other):
+        if isinstance(other, ABCSeries):
+            # the arrays defer to Series for comparison ops but the indexes
+            #  don't, so we have to unwrap here.
+            other = other._values
+
+        other = _maybe_unwrap_index(other)
+
+        op = getattr(self._data, opname)
+        return op(other)
+
+    wrapper.__name__ = opname
+    return wrapper
+
+
+def make_wrapped_arith_op(opname):
+    def method(self, other):
+        meth = getattr(self._data, opname)
+        result = meth(_maybe_unwrap_index(other))
+        return _wrap_arithmetic_op(self, other, result)
+
+    method.__name__ = opname
+    return method
+
+
+def _wrap_arithmetic_op(self, other, result):
+    if result is NotImplemented:
+        return NotImplemented
+
+    if isinstance(result, tuple):
+        # divmod, rdivmod
+        assert len(result) == 2
+        return (
+            _wrap_arithmetic_op(self, other, result[0]),
+            _wrap_arithmetic_op(self, other, result[1]),
+        )
+
+    if not isinstance(result, Index):
+        # Index.__new__ will choose appropriate subclass for dtype
+        result = Index(result)
+
+    res_name = get_op_result_name(self, other)
+    result.name = res_name
+    return result
+
+
+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, Index):
+        return obj._data
+    return obj

pandas/core/indexes/period.py

@@ -66,7 +66,6 @@ class PeriodDelegateMixin(DatetimelikeDelegateMixin):
     Delegate from PeriodIndex to PeriodArray.
     """
 
-    _delegate_class = PeriodArray
     _raw_methods = {"_format_native_types"}
     _raw_properties = {"is_leap_year", "freq"}
 

pandas/core/indexes/timedeltas.py

@@ -41,7 +41,6 @@ class TimedeltaDelegateMixin(DatetimelikeDelegateMixin):
     # Some are "raw" methods, the result is not re-boxed in an Index
     # We also have a few "extra" attrs, which may or may not be raw,
     # which we don't want to expose in the .dt accessor.
-    _delegate_class = TimedeltaArray
     _raw_properties = {"components", "_box_func"}
     _raw_methods = {"to_pytimedelta", "sum", "std", "median", "_format_native_types"}