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