Skip to content

Commit 09bd172

Browse files
jbrockmendeljreback
authored andcommitted
REF: move sharable methods to ExtensionIndex (#30717)
1 parent 1c3d64b commit 09bd172

File tree

3 files changed

+25
-23
lines changed

3 files changed

+25
-23
lines changed

pandas/core/indexes/datetimes.py

-14
Original file line numberDiff line numberDiff line change
@@ -830,20 +830,6 @@ def slice_indexer(self, start=None, end=None, step=None, kind=None):
830830
else:
831831
raise
832832

833-
# --------------------------------------------------------------------
834-
# Wrapping DatetimeArray
835-
836-
def __getitem__(self, key):
837-
result = self._data.__getitem__(key)
838-
if is_scalar(result):
839-
return result
840-
elif result.ndim > 1:
841-
# To support MPL which performs slicing with 2 dim
842-
# even though it only has 1 dim by definition
843-
assert isinstance(result, np.ndarray), result
844-
return result
845-
return type(self)(result, name=self.name)
846-
847833
# --------------------------------------------------------------------
848834

849835
@Substitution(klass="DatetimeIndex")

pandas/core/indexes/extension.py

+25
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
"""
44
from typing import List
55

6+
import numpy as np
7+
68
from pandas.compat.numpy import function as nv
79
from pandas.util._decorators import cache_readonly
810

@@ -170,6 +172,29 @@ class ExtensionIndex(Index):
170172
__le__ = _make_wrapped_comparison_op("__le__")
171173
__ge__ = _make_wrapped_comparison_op("__ge__")
172174

175+
def __getitem__(self, key):
176+
result = self._data[key]
177+
if isinstance(result, type(self._data)):
178+
return type(self)(result, name=self.name)
179+
180+
# Includes cases where we get a 2D ndarray back for MPL compat
181+
return result
182+
183+
def __iter__(self):
184+
return self._data.__iter__()
185+
186+
@property
187+
def _ndarray_values(self) -> np.ndarray:
188+
return self._data._ndarray_values
189+
190+
def dropna(self, how="any"):
191+
if how not in ("any", "all"):
192+
raise ValueError(f"invalid how option: {how}")
193+
194+
if self.hasnans:
195+
return self._shallow_copy(self._data[~self._isnan])
196+
return self._shallow_copy()
197+
173198
def repeat(self, repeats, axis=None):
174199
nv.validate_repeat(tuple(), dict(axis=axis))
175200
result = self._data.repeat(repeats, axis=axis)

pandas/core/indexes/timedeltas.py

-9
Original file line numberDiff line numberDiff line change
@@ -211,15 +211,6 @@ def _formatter_func(self):
211211

212212
return _get_format_timedelta64(self, box=True)
213213

214-
# -------------------------------------------------------------------
215-
# Wrapping TimedeltaArray
216-
217-
def __getitem__(self, key):
218-
result = self._data.__getitem__(key)
219-
if is_scalar(result):
220-
return result
221-
return type(self)(result, name=self.name)
222-
223214
# -------------------------------------------------------------------
224215

225216
@Appender(_index_shared_docs["astype"])

0 commit comments

Comments
 (0)