REF: move ExtensionIndex.map to be part of DatetimeLikeArrayMixin.map (#51934)

Author: topper-123

pandas/core/arrays/datetimelike.py

@@ -119,6 +119,7 @@
 from pandas.core.algorithms import (
  checked_add_with_arr,
  isin,
+ map_array,
  unique1d,
 )
 from pandas.core.array_algos import datetimelike_accumulations
@@ -754,14 +755,26 @@ def map(self, mapper, na_action=None):
  if na_action is not None:
  raise NotImplementedError
 
- # TODO(GH-23179): Add ExtensionArray.map
- # Need to figure out if we want ExtensionArray.map first.
- # If so, then we can refactor IndexOpsMixin._map_values to
- # a standalone function and call from here..
- # Else, just rewrite _map_infer_values to do the right thing.
  from pandas import Index
 
- return Index(self).map(mapper).array
+ idx = Index(self)
+ try:
+ result = mapper(idx)
+
+ # 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")
+ except Exception:
+ result = map_array(self, mapper)
+ result = Index(result)
+
+ if isinstance(result, ABCMultiIndex):
+ return result.to_numpy()
+ else:
+ return result.array
 
  def isin(self, values) -> npt.NDArray[np.bool_]:
  """

pandas/core/base.py

@@ -889,10 +889,8 @@ def _map_values(self, mapper, na_action=None):
  """
  arr = extract_array(self, extract_numpy=True, extract_range=True)
 
- if is_extension_array_dtype(arr.dtype):
- # Item "IndexOpsMixin" of "Union[IndexOpsMixin, ExtensionArray,
- # ndarray[Any, Any]]" has no attribute "map"
- return arr.map(mapper, na_action=na_action) # type: ignore[union-attr]
+ if isinstance(arr, ExtensionArray):
+ return arr.map(mapper, na_action=na_action)
 
  # Argument 1 to "map_array" has incompatible type
  # "Union[IndexOpsMixin, ExtensionArray, ndarray[Any, Any]]";

pandas/core/indexes/extension.py

@@ -9,18 +9,15 @@
  TypeVar,
 )
 
-import numpy as np
-
-from pandas.util._decorators import (
- cache_readonly,
- doc,
-)
+from pandas.util._decorators import cache_readonly
 
 from pandas.core.dtypes.generic import ABCDataFrame
 
 from pandas.core.indexes.base import Index
 
 if TYPE_CHECKING:
+ import numpy as np
+
  from pandas._typing import (
  ArrayLike,
  npt,
@@ -154,23 +151,6 @@ def _validate_fill_value(self, value):
  """
  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