""" Shared methods for Index subclasses backed by ExtensionArray. """ from __future__ import annotations from typing import ( TYPE_CHECKING, Callable, TypeVar, ) import numpy as np from pandas._typing import ( ArrayLike, npt, ) from pandas.util._decorators import ( cache_readonly, doc, ) from pandas.core.dtypes.generic import ABCDataFrame from pandas.core.indexes.base import Index if TYPE_CHECKING: from pandas.core.arrays import IntervalArray from pandas.core.arrays._mixins import NDArrayBackedExtensionArray _T = TypeVar("_T", bound="NDArrayBackedExtensionIndex") _ExtensionIndexT = TypeVar("_ExtensionIndexT", bound="ExtensionIndex") def _inherit_from_data( name: str, delegate: type, cache: bool = False, wrap: 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 wrap : bool, default False Whether to wrap the inherited result in an Index. Returns ------- attribute, method, property, or cache_readonly """ attr = getattr(delegate, name) if isinstance(attr, property) or type(attr).__name__ == "getset_descriptor": # getset_descriptor i.e. property defined in cython class if cache: def cached(self): return getattr(self._data, name) cached.__name__ = name cached.__doc__ = attr.__doc__ method = cache_readonly(cached) else: def fget(self): result = getattr(self._data, name) if wrap: if isinstance(result, type(self._data)): return type(self)._simple_new(result, name=self.name) elif isinstance(result, ABCDataFrame): return result.set_index(self) return Index(result, name=self.name) return result def fset(self, value) -> None: 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: # error: Incompatible redefinition (redefinition with type "Callable[[Any, # VarArg(Any), KwArg(Any)], Any]", original type "property") def method(self, *args, **kwargs): # type: ignore[misc] if "inplace" in kwargs: raise ValueError(f"cannot use inplace with {type(self).__name__}") result = attr(self._data, *args, **kwargs) if wrap: if isinstance(result, type(self._data)): return type(self)._simple_new(result, name=self.name) elif isinstance(result, ABCDataFrame): return result.set_index(self) return Index(result, name=self.name) return result # error: "property" has no attribute "__name__" method.__name__ = name # type: ignore[attr-defined] method.__doc__ = attr.__doc__ return method def inherit_names( names: list[str], delegate: type, cache: bool = False, wrap: bool = False ) -> Callable[[type[_ExtensionIndexT]], type[_ExtensionIndexT]]: """ Class decorator to pin attributes from an ExtensionArray to a Index subclass. Parameters ---------- names : List[str] delegate : class cache : bool, default False wrap : bool, default False Whether to wrap the inherited result in an Index. """ def wrapper(cls: type[_ExtensionIndexT]) -> type[_ExtensionIndexT]: for name in names: meth = _inherit_from_data(name, delegate, cache=cache, wrap=wrap) setattr(cls, name, meth) return cls return wrapper class ExtensionIndex(Index): """ Index subclass for indexes backed by ExtensionArray. """ # The base class already passes through to _data: # size, __len__, dtype _data: IntervalArray | NDArrayBackedExtensionArray # --------------------------------------------------------------------- def _validate_fill_value(self, value): """ Convert value to be insertable to underlying array. """ return self._data._validate_setitem_value(value) @doc(Index.map) def map(self, mapper, na_action=None): # Try to run function on index first, and then on elements of index # Especially important for group-by functionality try: result = mapper(self) # Try to use this result if we can if isinstance(result, np.ndarray): result = Index(result) if not isinstance(result, Index): raise TypeError("The map function must return an Index object") return result except Exception: return self.astype(object).map(mapper) @cache_readonly def _isnan(self) -> npt.NDArray[np.bool_]: # error: Incompatible return value type (got "ExtensionArray", expected # "ndarray") return self._data.isna() # type: ignore[return-value] class NDArrayBackedExtensionIndex(ExtensionIndex): """ Index subclass for indexes backed by NDArrayBackedExtensionArray. """ _data: NDArrayBackedExtensionArray def _get_engine_target(self) -> np.ndarray: return self._data._ndarray def _from_join_target(self, result: np.ndarray) -> ArrayLike: assert result.dtype == self._data._ndarray.dtype return self._data._from_backing_data(result)