123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192 |
- """
- 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)
|