diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index 1bfec9fbad0ed..e50d752f910e0 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -164,22 +164,6 @@ def __contains__(self, key: Any) -> bool: is_scalar(res) or isinstance(res, slice) or (is_list_like(res) and len(res)) ) - # 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. diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index db35cdb72979f..9ddc5c01030b1 100644 --- a/pandas/core/indexes/extension.py +++ b/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: diff --git a/pandas/tests/indexes/categorical/test_category.py b/pandas/tests/indexes/categorical/test_category.py index e027641288bb9..c69c1f3f386a9 100644 --- a/pandas/tests/indexes/categorical/test_category.py +++ b/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) diff --git a/pandas/tests/indexes/common.py b/pandas/tests/indexes/common.py index f3ebe8313d0c6..74731fa13b629 100644 --- a/pandas/tests/indexes/common.py +++ b/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()