REF: implement indexes.extension to share delegation (#30629)

Author: jbrockmendel

pandas/core/indexes/datetimelike.py

@@ -2,7 +2,7 @@
 Base and utility classes for tseries type pandas objects.
 """
 import operator
-from typing import List, Set
+from typing import List, Optional, Set
 
 import numpy as np
 
@@ -40,28 +40,9 @@
 
 from pandas.tseries.frequencies import DateOffset, to_offset
 
-_index_doc_kwargs = dict(ibase._index_doc_kwargs)
-
+from .extension import inherit_names
 
-def ea_passthrough(array_method):
-    """
-    Make an alias for a method of the underlying ExtensionArray.
-
-    Parameters
-    ----------
-    array_method : method on an Array class
-
-    Returns
-    -------
-    method
-    """
-
-    def method(self, *args, **kwargs):
-        return array_method(self._data, *args, **kwargs)
-
-    method.__name__ = array_method.__name__
-    method.__doc__ = array_method.__doc__
-    return method
+_index_doc_kwargs = dict(ibase._index_doc_kwargs)
 
 
 def _make_wrapped_arith_op(opname):
@@ -100,48 +81,34 @@ def wrapper(left, right):
     return wrapper
 
 
+@inherit_names(
+    ["inferred_freq", "_isnan", "_resolution", "resolution"],
+    DatetimeLikeArrayMixin,
+    cache=True,
+)
+@inherit_names(
+    ["__iter__", "mean", "freq", "freqstr", "_ndarray_values", "asi8", "_box_values"],
+    DatetimeLikeArrayMixin,
+)
 class DatetimeIndexOpsMixin(ExtensionOpsMixin):
     """
     Common ops mixin to support a unified interface datetimelike Index.
     """
 
     _data: ExtensionArray
+    freq: Optional[DateOffset]
+    freqstr: Optional[str]
+    _resolution: int
+    _bool_ops: List[str] = []
+    _field_ops: List[str] = []
 
-    # DatetimeLikeArrayMixin assumes subclasses are mutable, so these are
-    # properties there.  They can be made into cache_readonly for Index
-    # subclasses bc they are immutable
-    inferred_freq = cache_readonly(
-        DatetimeLikeArrayMixin.inferred_freq.fget  # type: ignore
-    )
-    _isnan = cache_readonly(DatetimeLikeArrayMixin._isnan.fget)  # type: ignore
     hasnans = cache_readonly(DatetimeLikeArrayMixin._hasnans.fget)  # type: ignore
     _hasnans = hasnans  # for index / array -agnostic code
-    _resolution = cache_readonly(
-        DatetimeLikeArrayMixin._resolution.fget  # type: ignore
-    )
-    resolution = cache_readonly(DatetimeLikeArrayMixin.resolution.fget)  # type: ignore
-
-    __iter__ = ea_passthrough(DatetimeLikeArrayMixin.__iter__)
-    mean = ea_passthrough(DatetimeLikeArrayMixin.mean)
 
     @property
     def is_all_dates(self) -> bool:
         return True
 
-    @property
-    def freq(self):
-        """
-        Return the frequency object if it is set, otherwise None.
-        """
-        return self._data.freq
-
-    @property
-    def freqstr(self):
-        """
-        Return the frequency object as a string if it is set, otherwise None.
-        """
-        return self._data.freqstr
-
     def unique(self, level=None):
         if level is not None:
             self._validate_index_level(level)
@@ -172,10 +139,6 @@ def wrapper(self, other):
         wrapper.__name__ = f"__{op.__name__}__"
         return wrapper
 
-    @property
-    def _ndarray_values(self) -> np.ndarray:
-        return self._data._ndarray_values
-
     # ------------------------------------------------------------------------
     # Abstract data attributes
 
@@ -184,11 +147,6 @@ def values(self):
         # Note: PeriodArray overrides this to return an ndarray of objects.
         return self._data._data
 
-    @property  # type: ignore # https://github.com/python/mypy/issues/1362
-    @Appender(DatetimeLikeArrayMixin.asi8.__doc__)
-    def asi8(self):
-        return self._data.asi8
-
     def __array_wrap__(self, result, context=None):
         """
         Gets called after a ufunc.
@@ -248,9 +206,6 @@ def _ensure_localized(
             return type(self)._simple_new(result, name=self.name)
         return arg
 
-    def _box_values(self, values):
-        return self._data._box_values(values)
-
     @Appender(_index_shared_docs["contains"] % _index_doc_kwargs)
     def __contains__(self, key):
         try:

pandas/core/indexes/datetimes.py

@@ -35,6 +35,7 @@
     DatetimelikeDelegateMixin,
     DatetimeTimedeltaMixin,
 )
+from pandas.core.indexes.extension import inherit_names
 from pandas.core.ops import get_op_result_name
 import pandas.core.tools.datetimes as tools
 
@@ -72,6 +73,7 @@ class DatetimeDelegateMixin(DatetimelikeDelegateMixin):
         "_local_timestamps",
         "_has_same_tz",
         "_format_native_types",
+        "__iter__",
     ]
     _extra_raw_properties = ["_box_func", "tz", "tzinfo", "dtype"]
     _delegated_properties = DatetimeArray._datetimelike_ops + _extra_raw_properties
@@ -87,6 +89,17 @@ class DatetimeDelegateMixin(DatetimelikeDelegateMixin):
     _delegate_class = DatetimeArray
 
 
+@inherit_names(["_timezone", "is_normalized", "_resolution"], DatetimeArray, cache=True)
+@inherit_names(
+    [
+        "_bool_ops",
+        "_object_ops",
+        "_field_ops",
+        "_datetimelike_ops",
+        "_datetimelike_methods",
+    ],
+    DatetimeArray,
+)
 @delegate_names(
     DatetimeArray, DatetimeDelegateMixin._delegated_properties, typ="property"
 )
@@ -209,15 +222,6 @@ class DatetimeIndex(DatetimeTimedeltaMixin, DatetimeDelegateMixin):
     _is_numeric_dtype = False
     _infer_as_myclass = True
 
-    # Use faster implementation given we know we have DatetimeArrays
-    __iter__ = DatetimeArray.__iter__
-    # some things like freq inference make use of these attributes.
-    _bool_ops = DatetimeArray._bool_ops
-    _object_ops = DatetimeArray._object_ops
-    _field_ops = DatetimeArray._field_ops
-    _datetimelike_ops = DatetimeArray._datetimelike_ops
-    _datetimelike_methods = DatetimeArray._datetimelike_methods
-
     tz: Optional[tzinfo]
 
     # --------------------------------------------------------------------
@@ -947,10 +951,6 @@ def slice_indexer(self, start=None, end=None, step=None, kind=None):
     # --------------------------------------------------------------------
     # Wrapping DatetimeArray
 
-    _timezone = cache_readonly(DatetimeArray._timezone.fget)  # type: ignore
-    is_normalized = cache_readonly(DatetimeArray.is_normalized.fget)  # type: ignore
-    _resolution = cache_readonly(DatetimeArray._resolution.fget)  # type: ignore
-
     def __getitem__(self, key):
         result = self._data.__getitem__(key)
         if is_scalar(result):

pandas/core/indexes/extension.py

@@ -0,0 +1,78 @@
+"""
+Shared methods for Index subclasses backed by ExtensionArray.
+"""
+from typing import List
+
+from pandas.util._decorators import cache_readonly
+
+
+def inherit_from_data(name: str, delegate, cache: bool = False):
+    """
+    Make an alias for a method of the underlying ExtensionArray.
+
+    Parameters
+    ----------
+    name : str
+        Name of an attribute the class should inherit from its EA parent.
+    delegate : class
+    cache : bool, default False
+        Whether to convert wrapped properties into cache_readonly
+
+    Returns
+    -------
+    attribute, method, property, or cache_readonly
+    """
+
+    attr = getattr(delegate, name)
+
+    if isinstance(attr, property):
+        if cache:
+            method = cache_readonly(attr.fget)
+
+        else:
+
+            def fget(self):
+                return getattr(self._data, name)
+
+            def fset(self, value):
+                setattr(self._data, name, value)
+
+            fget.__name__ = name
+            fget.__doc__ = attr.__doc__
+
+            method = property(fget, fset)
+
+    elif not callable(attr):
+        # just a normal attribute, no wrapping
+        method = attr
+
+    else:
+
+        def method(self, *args, **kwargs):
+            result = attr(self._data, *args, **kwargs)
+            return result
+
+        method.__name__ = name
+        method.__doc__ = attr.__doc__
+    return method
+
+
+def inherit_names(names: List[str], delegate, cache: bool = False):
+    """
+    Class decorator to pin attributes from an ExtensionArray to a Index subclass.
+
+    Parameters
+    ----------
+    names : List[str]
+    delegate : class
+    cache : bool, default False
+    """
+
+    def wrapper(cls):
+        for name in names:
+            meth = inherit_from_data(name, delegate, cache=cache)
+            setattr(cls, name, meth)
+
+        return cls
+
+    return wrapper

pandas/core/indexes/interval.py

@@ -58,6 +58,8 @@
 from pandas.tseries.frequencies import to_offset
 from pandas.tseries.offsets import DateOffset
 
+from .extension import inherit_names
+
 _VALID_CLOSED = {"left", "right", "both", "neither"}
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
 
@@ -199,10 +201,11 @@ def func(intvidx_self, other, sort=False):
 )
 @accessor.delegate_names(
     delegate=IntervalArray,
-    accessors=["__array__", "overlaps", "contains"],
+    accessors=["__array__", "overlaps", "contains", "__len__", "set_closed"],
     typ="method",
     overwrite=True,
 )
+@inherit_names(["is_non_overlapping_monotonic", "mid"], IntervalArray, cache=True)
 class IntervalIndex(IntervalMixin, Index, accessor.PandasDelegate):
     _typ = "intervalindex"
     _comparables = ["name"]
@@ -412,34 +415,6 @@ def to_tuples(self, na_tuple=True):
     def _multiindex(self):
         return MultiIndex.from_arrays([self.left, self.right], names=["left", "right"])
 
-    @Appender(
-        _interval_shared_docs["set_closed"]
-        % dict(
-            klass="IntervalIndex",
-            examples=textwrap.dedent(
-                """\
-        Examples
-        --------
-        >>> index = pd.interval_range(0, 3)
-        >>> index
-        IntervalIndex([(0, 1], (1, 2], (2, 3]],
-                      closed='right',
-                      dtype='interval[int64]')
-        >>> index.set_closed('both')
-        IntervalIndex([[0, 1], [1, 2], [2, 3]],
-                      closed='both',
-                      dtype='interval[int64]')
-        """
-            ),
-        )
-    )
-    def set_closed(self, closed):
-        array = self._data.set_closed(closed)
-        return self._simple_new(array, self.name)  # TODO: can we use _shallow_copy?
-
-    def __len__(self) -> int:
-        return len(self.left)
-
     @cache_readonly
     def values(self):
         """
@@ -479,13 +454,6 @@ def memory_usage(self, deep: bool = False) -> int:
         # so return the bytes here
         return self.left.memory_usage(deep=deep) + self.right.memory_usage(deep=deep)
 
-    @cache_readonly
-    def mid(self):
-        """
-        Return the midpoint of each Interval in the IntervalIndex as an Index.
-        """
-        return self._data.mid
-
     @cache_readonly
     def is_monotonic(self) -> bool:
         """
@@ -534,11 +502,6 @@ def is_unique(self):
 
         return True
 
-    @cache_readonly
-    @Appender(_interval_shared_docs["is_non_overlapping_monotonic"] % _index_doc_kwargs)
-    def is_non_overlapping_monotonic(self):
-        return self._data.is_non_overlapping_monotonic
-
     @property
     def is_overlapping(self):
         """