diff --git a/pandas/core/indexes/category.py b/pandas/core/indexes/category.py index c162bbb412591..f61721a0e51e6 100644 --- a/pandas/core/indexes/category.py +++ b/pandas/core/indexes/category.py @@ -10,7 +10,6 @@ from pandas._libs.hashtable import duplicated_int64 from pandas._typing import AnyArrayLike import pandas.compat as compat -from pandas.compat.numpy import function as nv from pandas.util._decorators import Appender, cache_readonly from pandas.core.dtypes.common import ( @@ -30,7 +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 +from pandas.core.indexes.extension import ExtensionIndex, make_wrapped_comparison_op import pandas.core.missing as missing from pandas.core.ops import get_op_result_name @@ -67,7 +66,7 @@ typ="method", overwrite=True, ) -class CategoricalIndex(Index, accessor.PandasDelegate): +class CategoricalIndex(ExtensionIndex, accessor.PandasDelegate): """ Index based on an underlying :class:`Categorical`. @@ -723,20 +722,8 @@ def _convert_arr_indexer(self, keyarr): def _convert_index_indexer(self, keyarr): return self._shallow_copy(keyarr) - @Appender(_index_shared_docs["take"] % _index_doc_kwargs) - def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs): - nv.validate_take(tuple(), kwargs) - indices = ensure_platform_int(indices) - taken = self._assert_take_fillable( - self._data, - indices, - allow_fill=allow_fill, - fill_value=fill_value, - na_value=self._data.dtype.na_value, - ) - return self._shallow_copy(taken) - def take_nd(self, *args, **kwargs): + """Alias for `take`""" warnings.warn( "CategoricalIndex.take_nd is deprecated, use CategoricalIndex.take instead", FutureWarning, diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index 365cfb1585f00..0c5f451c4f07b 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -41,7 +41,12 @@ from pandas.tseries.frequencies import DateOffset, to_offset -from .extension import inherit_names, make_wrapped_arith_op, make_wrapped_comparison_op +from .extension import ( + ExtensionIndex, + inherit_names, + make_wrapped_arith_op, + make_wrapped_comparison_op, +) _index_doc_kwargs = dict(ibase._index_doc_kwargs) @@ -81,7 +86,7 @@ def wrapper(left, right): ["__iter__", "mean", "freq", "freqstr", "_ndarray_values", "asi8", "_box_values"], DatetimeLikeArrayMixin, ) -class DatetimeIndexOpsMixin(ExtensionOpsMixin): +class DatetimeIndexOpsMixin(ExtensionIndex, ExtensionOpsMixin): """ Common ops mixin to support a unified interface datetimelike Index. """ @@ -246,16 +251,13 @@ def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs): if isinstance(maybe_slice, slice): return self[maybe_slice] - taken = self._assert_take_fillable( - self._data, - indices, - allow_fill=allow_fill, - fill_value=fill_value, - na_value=NaT, + taken = ExtensionIndex.take( + self, indices, axis, allow_fill, fill_value, **kwargs ) # keep freq in PeriodArray/Index, reset otherwise freq = self.freq if is_period_dtype(self) else None + assert taken.freq == freq, (taken.freq, freq, taken) return self._shallow_copy(taken, freq=freq) _can_hold_na = True diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index 3c98d31e34b7d..6f581c4ebb594 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -3,10 +3,13 @@ """ from typing import List +from pandas.compat.numpy import function as nv from pandas.util._decorators import cache_readonly +from pandas.core.dtypes.common import ensure_platform_int from pandas.core.dtypes.generic import ABCSeries +from pandas.core.arrays import ExtensionArray from pandas.core.ops import get_op_result_name from .base import Index @@ -152,3 +155,24 @@ def _maybe_unwrap_index(obj): if isinstance(obj, Index): return obj._data return obj + + +class ExtensionIndex(Index): + """ + Index subclass for indexes backed by ExtensionArray. + """ + + _data: ExtensionArray + + def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs): + nv.validate_take(tuple(), kwargs) + indices = ensure_platform_int(indices) + + taken = self._assert_take_fillable( + self._data, + indices, + allow_fill=allow_fill, + fill_value=fill_value, + na_value=self._na_value, + ) + return type(self)(taken, name=self.name) diff --git a/pandas/core/indexes/interval.py b/pandas/core/indexes/interval.py index 35cdf840a55b2..2677ed4d95158 100644 --- a/pandas/core/indexes/interval.py +++ b/pandas/core/indexes/interval.py @@ -58,7 +58,7 @@ from pandas.tseries.frequencies import to_offset from pandas.tseries.offsets import DateOffset -from .extension import inherit_names +from .extension import ExtensionIndex, inherit_names _VALID_CLOSED = {"left", "right", "both", "neither"} _index_doc_kwargs = dict(ibase._index_doc_kwargs) @@ -213,7 +213,7 @@ def func(intvidx_self, other, sort=False): overwrite=True, ) @inherit_names(["is_non_overlapping_monotonic", "mid"], IntervalArray, cache=True) -class IntervalIndex(IntervalMixin, Index, accessor.PandasDelegate): +class IntervalIndex(IntervalMixin, ExtensionIndex, accessor.PandasDelegate): _typ = "intervalindex" _comparables = ["name"] _attributes = ["name", "closed"]