From e78d2f2a0da3451d8aea2c04338944627e05b9d0 Mon Sep 17 00:00:00 2001
From: Terji Petersen <terji78@gmail.com>
Date: Mon, 13 Mar 2023 18:25:41 +0000
Subject: [PATCH] REF: move ExtensionIndex.map to be part of
 DatetimeLikeArrayMixin.map

---
 pandas/core/arrays/datetimelike.py | 25 +++++++++++++++++++------
 pandas/core/base.py                |  6 ++----
 pandas/core/indexes/extension.py   | 26 +++-----------------------
 3 files changed, 24 insertions(+), 33 deletions(-)

diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py
index 5295f89626e04..b1e2cdd104937 100644
--- a/pandas/core/arrays/datetimelike.py
+++ b/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_]:
         """
diff --git a/pandas/core/base.py b/pandas/core/base.py
index 2d3338d7c9f10..6218fd3e9931e 100644
--- a/pandas/core/base.py
+++ b/pandas/core/base.py
@@ -881,10 +881,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]]";
diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py
index 8c8de96ad5b54..7d4dcf54a025f 100644
--- a/pandas/core/indexes/extension.py
+++ b/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