Backport PR pandas-dev#31232: REGR: Fix IntervalIndex.map when result is object dtype (pandas-dev#31246)

Co-authored-by: Jeremy Schendel <[email protected]>

Author: 2 people

pandas/core/indexes/datetimelike.py

@@ -164,22 +164,6 @@ def __contains__(self, key):
         except (KeyError, TypeError, ValueError):
             return False
 
-    # Try to run function on index first, and then on elements of index
-    # Especially important for group-by functionality
-    def map(self, mapper, na_action=None):
-        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)
-
     def sort_values(self, return_indexer=False, ascending=True):
         """
         Return sorted copy of Index.

pandas/core/indexes/extension.py

@@ -232,6 +232,23 @@ def _get_unique_index(self, dropna=False):
             result = result[~result.isna()]
         return self._shallow_copy(result)
 
+    @Appender(Index.map.__doc__)
+    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)
+
     @Appender(Index.astype.__doc__)
     def astype(self, dtype, copy=True):
         if is_dtype_equal(self.dtype, dtype) and copy is False:

pandas/tests/indexes/categorical/test_category.py

@@ -981,3 +981,20 @@ def test_getitem_2d_deprecated(self):
         idx = self.create_index()
         with pytest.raises(ValueError, match="cannot mask with array containing NA"):
             idx[:, None]
+
+    @pytest.mark.parametrize(
+        "data, categories",
+        [
+            (list("abcbca"), list("cab")),
+            (pd.interval_range(0, 3).repeat(3), pd.interval_range(0, 3)),
+        ],
+        ids=["string", "interval"],
+    )
+    def test_map_str(self, data, categories, ordered_fixture):
+        # GH 31202 - override base class since we want to maintain categorical/ordered
+        index = CategoricalIndex(data, categories=categories, ordered=ordered_fixture)
+        result = index.map(str)
+        expected = CategoricalIndex(
+            map(str, data), categories=map(str, categories), ordered=ordered_fixture
+        )
+        tm.assert_index_equal(result, expected)

pandas/tests/indexes/common.py

@@ -808,6 +808,13 @@ def test_map_dictlike(self, mapper):
         result = index.map(mapper(expected, index))
         tm.assert_index_equal(result, expected)
 
+    def test_map_str(self):
+        # GH 31202
+        index = self.create_index()
+        result = index.map(str)
+        expected = Index([str(x) for x in index], dtype=object)
+        tm.assert_index_equal(result, expected)
+
     def test_putmask_with_wrong_mask(self):
         # GH18368
         index = self.create_index()