REF: consolidate DTI/TDI/PI get_value in ExtensionIndex, associated cleanups (#31406)

Author: jbrockmendel

pandas/core/indexes/base.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 import operator
 from textwrap import dedent
-from typing import Any, FrozenSet, Hashable, Optional, Union
+from typing import TYPE_CHECKING, Any, FrozenSet, Hashable, Optional, Union
 import warnings
 
 import numpy as np
@@ -83,6 +83,10 @@
     pprint_thing,
 )
 
+if TYPE_CHECKING:
+    from pandas import Series
+
+
 __all__ = ["Index"]
 
 _unsortable_types = frozenset(("mixed", "mixed-integer"))
@@ -4577,21 +4581,15 @@ def argsort(self, *args, **kwargs) -> np.ndarray:
             result = np.array(self)
         return result.argsort(*args, **kwargs)
 
-    _index_shared_docs[
-        "get_value"
-    ] = """
+    def get_value(self, series: "Series", key):
+        """
         Fast lookup of value from 1-dimensional ndarray. Only use this if you
         know what you're doing.
 
         Returns
         -------
-        scalar
-            A value in the Series with the index of the key value in self.
+        scalar or Series
         """
-
-    @Appender(_index_shared_docs["get_value"] % _index_doc_kwargs)
-    def get_value(self, series, key):
-
         if not is_scalar(key):
             # if key is not a scalar, directly raise an error (the code below
             # would convert to numpy arrays and raise later any way) - GH29926
@@ -4616,7 +4614,7 @@ def get_value(self, series, key):
 
         return self._get_values_for_loc(series, loc)
 
-    def _get_values_for_loc(self, series, loc):
+    def _get_values_for_loc(self, series: "Series", loc):
         """
         Do a positional lookup on the given Series, returning either a scalar
         or a Series.

pandas/core/indexes/category.py

@@ -1,4 +1,4 @@
-from typing import Any, List
+from typing import TYPE_CHECKING, Any, List
 import warnings
 
 import numpy as np
@@ -7,7 +7,6 @@
 
 from pandas._libs import index as libindex
 from pandas._libs.hashtable import duplicated_int64
-from pandas._typing import AnyArrayLike
 from pandas.util._decorators import Appender, cache_readonly
 
 from pandas.core.dtypes.common import (
@@ -31,6 +30,9 @@
 import pandas.core.missing as missing
 from pandas.core.ops import get_op_result_name
 
+if TYPE_CHECKING:
+    from pandas import Series
+
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
 _index_doc_kwargs.update(dict(target_klass="CategoricalIndex"))
 
@@ -494,14 +496,14 @@ def get_loc(self, key, method=None):
         except KeyError:
             raise KeyError(key)
 
-    def get_value(self, series: AnyArrayLike, key: Any):
+    def get_value(self, series: "Series", key: Any):
         """
         Fast lookup of value from 1-dimensional ndarray. Only use this if you
         know what you're doing
 
         Parameters
         ----------
-        series : Series, ExtensionArray, Index, or ndarray
+        series : Series
             1-dimensional array to take values from
         key: : scalar
             The value of this index at the position of the desired value,
@@ -521,7 +523,7 @@ def get_value(self, series: AnyArrayLike, key: Any):
             pass
 
         # we might be a positional inexer
-        return super().get_value(series, key)
+        return Index.get_value(self, series, key)
 
     @Appender(Index.where.__doc__)
     def where(self, cond, other=None):

pandas/core/indexes/datetimelike.py

@@ -383,6 +383,7 @@ def _format_attrs(self):
         return attrs
 
     # --------------------------------------------------------------------
+    # Indexing Methods
 
     def _convert_scalar_indexer(self, key, kind=None):
         """
@@ -409,6 +410,8 @@ def _convert_scalar_indexer(self, key, kind=None):
 
         return super()._convert_scalar_indexer(key, kind=kind)
 
+    # --------------------------------------------------------------------
+
     __add__ = make_wrapped_arith_op("__add__")
     __radd__ = make_wrapped_arith_op("__radd__")
     __sub__ = make_wrapped_arith_op("__sub__")

pandas/core/indexes/datetimes.py

@@ -616,17 +616,6 @@ def _maybe_promote(self, other):
             other = DatetimeIndex(other)
         return self, other
 
-    def get_value(self, series, key):
-        """
-        Fast lookup of value from 1-dimensional ndarray. Only use this if you
-        know what you're doing
-        """
-        if is_integer(key):
-            loc = key
-        else:
-            loc = self.get_loc(key)
-        return self._get_values_for_loc(series, loc)
-
     def get_loc(self, key, method=None, tolerance=None):
         """
         Get integer location for requested label
@@ -642,7 +631,7 @@ def get_loc(self, key, method=None, tolerance=None):
         if is_valid_nat_for_dtype(key, self.dtype):
             key = NaT
 
-        if isinstance(key, (datetime, np.datetime64)):
+        if isinstance(key, self._data._recognized_scalars):
             # needed to localize naive datetimes
             key = self._maybe_cast_for_get_loc(key)
 

pandas/core/indexes/extension.py

@@ -1,7 +1,7 @@
 """
 Shared methods for Index subclasses backed by ExtensionArray.
 """
-from typing import List
+from typing import TYPE_CHECKING, List
 
 import numpy as np
 
@@ -11,6 +11,7 @@
 from pandas.core.dtypes.common import (
     ensure_platform_int,
     is_dtype_equal,
+    is_integer,
     is_object_dtype,
 )
 from pandas.core.dtypes.generic import ABCSeries
@@ -20,6 +21,9 @@
 from pandas.core.indexes.base import Index
 from pandas.core.ops import get_op_result_name
 
+if TYPE_CHECKING:
+    from pandas import Series
+
 
 def inherit_from_data(name: str, delegate, cache: bool = False, wrap: bool = False):
     """
@@ -293,3 +297,26 @@ def astype(self, dtype, copy=True):
         # pass copy=False because any copying will be done in the
         #  _data.astype call above
         return Index(new_values, dtype=new_values.dtype, name=self.name, copy=False)
+
+    # --------------------------------------------------------------------
+    # Indexing Methods
+
+    @Appender(Index.get_value.__doc__)
+    def get_value(self, series: "Series", key):
+        """
+        Fast lookup of value from 1-dimensional ndarray. Only use this if you
+        know what you're doing
+        """
+        try:
+            loc = self.get_loc(key)
+        except KeyError:
+            # e.g. DatetimeIndex doesn't hold integers
+            if is_integer(key) and not self.holds_integer():
+                # Fall back to positional
+                loc = key
+            else:
+                raise
+
+        return self._get_values_for_loc(series, loc)
+
+    # --------------------------------------------------------------------

pandas/core/indexes/interval.py

@@ -1,7 +1,7 @@
 """ define the IntervalIndex """
 from operator import le, lt
 import textwrap
-from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import numpy as np
 
@@ -57,10 +57,6 @@
 from pandas.tseries.frequencies import to_offset
 from pandas.tseries.offsets import DateOffset
 
-if TYPE_CHECKING:
-    from pandas import Series
-
-
 _VALID_CLOSED = {"left", "right", "both", "neither"}
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
 
@@ -527,6 +523,10 @@ def is_overlapping(self) -> bool:
         # GH 23309
         return self._engine.is_overlapping
 
+    def holds_integer(self):
+        return self.dtype.subtype.kind not in ["m", "M"]
+        # TODO: There must already exist something for this?
+
     @Appender(Index._convert_scalar_indexer.__doc__)
     def _convert_scalar_indexer(self, key, kind=None):
         if kind == "iloc":
@@ -884,11 +884,6 @@ def get_indexer_for(self, target: AnyArrayLike, **kwargs) -> np.ndarray:
             return self.get_indexer_non_unique(target)[0]
         return self.get_indexer(target, **kwargs)
 
-    @Appender(_index_shared_docs["get_value"] % _index_doc_kwargs)
-    def get_value(self, series: "Series", key):
-        loc = self.get_loc(key)
-        return series.iloc[loc]
-
     def _convert_slice_indexer(self, key: slice, kind=None):
         if not (key.step is None or key.step == 1):
             raise ValueError("cannot support not-default step in a slice")

pandas/core/indexes/numeric.py

@@ -416,6 +416,7 @@ def _format_native_types(
         )
         return formatter.get_result_as_array()
 
+    @Appender(Index.get_value.__doc__)
     def get_value(self, series: "Series", key):
         """
         We always want to get an index value, never a value.

pandas/core/indexes/period.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any
+from typing import Any
 import weakref
 
 import numpy as np
@@ -18,7 +18,6 @@
     is_float,
     is_integer,
     is_integer_dtype,
-    is_list_like,
     is_object_dtype,
     is_scalar,
     pandas_dtype,
@@ -51,9 +50,6 @@
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
 _index_doc_kwargs.update(dict(target_klass="PeriodIndex or list of Periods"))
 
-if TYPE_CHECKING:
-    from pandas import Series
-
 # --- Period index sketch
 
 
@@ -455,17 +451,6 @@ def inferred_type(self) -> str:
         # indexing
         return "period"
 
-    def get_value(self, series: "Series", key):
-        """
-        Fast lookup of value from 1-dimensional ndarray. Only use this if you
-        know what you're doing
-        """
-        if is_integer(key):
-            loc = key
-        else:
-            loc = self.get_loc(key)
-        return self._get_values_for_loc(series, loc)
-
     @Appender(_index_shared_docs["get_indexer"] % _index_doc_kwargs)
     def get_indexer(self, target, method=None, limit=None, tolerance=None):
         target = ensure_index(target)
@@ -560,9 +545,6 @@ def get_loc(self, key, method=None, tolerance=None):
             key = Period(key, freq=self.freq)
         except ValueError:
             # we cannot construct the Period
-            # as we have an invalid type
-            if is_list_like(key):
-                raise TypeError(f"'{key}' is an invalid key")
             raise KeyError(key)
 
         ordinal = key.ordinal if key is not NaT else key.value

pandas/core/indexes/timedeltas.py

@@ -222,17 +222,6 @@ def _maybe_promote(self, other):
             other = TimedeltaIndex(other)
         return self, other
 
-    def get_value(self, series, key):
-        """
-        Fast lookup of value from 1-dimensional ndarray. Only use this if you
-        know what you're doing
-        """
-        if is_integer(key):
-            loc = key
-        else:
-            loc = self.get_loc(key)
-        return self._get_values_for_loc(series, loc)
-
     def get_loc(self, key, method=None, tolerance=None):
         """
         Get integer location for requested label