make map return an index if it operates on an index, multi index, or categorical index; map on a categorical will either return a categorical or an index (rather than a numpy array)

Author: nateyoder

pandas/core/categorical.py

@@ -943,7 +943,7 @@ def map(self, mapper):
 
         Returns
         -------
-        applied : Categorical or np.ndarray.
+        applied : Categorical or Index.
         """
         new_categories = self.categories.map(mapper)
         try:

pandas/indexes/base.py

@@ -2428,7 +2428,7 @@ def groupby(self, values):
 
     def map(self, mapper):
         """
-        Apply mapper function to its values.
+        Apply mapper function to an index.
 
         Parameters
         ----------
@@ -2437,9 +2437,12 @@ def map(self, mapper):
 
         Returns
         -------
-        applied : array
+        applied : Index
+            The output of the mapping function applied to the index.
         """
-        return self._arrmap(self.values, mapper)
+        attributes = self._get_attributes_dict()
+        attributes['copy'] = False
+        return Index(self._arrmap(self.values, mapper), **attributes)
 
     def isin(self, values, level=None):
         """

pandas/indexes/category.py

@@ -524,15 +524,15 @@ def map(self, mapper):
         ----------
         mapper : callable
             Function to be applied. When all categories are mapped
-            to different categories, the result will be Categorical which has
-            the same order property as the original. Otherwise, the result will
-            be np.ndarray.
+            to different categories, the result will be a CategoricalIndex
+            which has the same order property as the original. Otherwise,
+            the result will be a Index.
 
         Returns
         -------
-        applied : Categorical or np.ndarray.
+        applied : CategoricalIndex or Index
         """
-        return self.values.map(mapper)
+        return self._shallow_copy_with_infer(self.values.map(mapper))
 
     def delete(self, loc):
         """

pandas/tests/indexes/test_base.py

@@ -767,6 +767,26 @@ def test_sub(self):
         self.assertRaises(TypeError, lambda: idx - idx.tolist())
         self.assertRaises(TypeError, lambda: idx.tolist() - idx)
 
+    def test_map_identity_mapping(self):
+        for name, cur_index in self.indices.items():
+            if name == 'tuples':
+                expected = Index(cur_index.values, tupleize_cols=False)
+                self.assert_index_equal(expected, cur_index.map(lambda x: x))
+            else:
+                self.assert_index_equal(cur_index, cur_index.map(lambda x: x))
+
+    def test_map_that_returns_tuples_creates_index_not_multi_index(self):
+        boolean_index = tm.makeIntIndex(3).map(lambda x: (x, x == 1))
+        expected = Index([(0, False), (1, True), (2, False)],
+                         tupleize_cols=False)
+        self.assert_index_equal(boolean_index, expected)
+
+    def test_map_that_reduces_multi_index_to_single_index_returns_index(self):
+        first_level = ['foo', 'bar', 'baz']
+        multi_index = MultiIndex.from_tuples(lzip(first_level, [1, 2, 3]))
+        reduced_index = multi_index.map(lambda x: x[0])
+        self.assert_index_equal(reduced_index, Index(first_level))
+
     def test_append_multiple(self):
         index = Index(['a', 'b', 'c', 'd', 'e', 'f'])
 
@@ -1194,16 +1214,16 @@ def check_slice(in_slice, expected):
             self.assert_index_equal(result, expected)
 
         for in_slice, expected in [
-                (SLC[::-1], 'yxdcb'), (SLC['b':'y':-1], ''),
-                (SLC['b'::-1], 'b'), (SLC[:'b':-1], 'yxdcb'),
-                (SLC[:'y':-1], 'y'), (SLC['y'::-1], 'yxdcb'),
-                (SLC['y'::-4], 'yb'),
-                # absent labels
-                (SLC[:'a':-1], 'yxdcb'), (SLC[:'a':-2], 'ydb'),
-                (SLC['z'::-1], 'yxdcb'), (SLC['z'::-3], 'yc'),
-                (SLC['m'::-1], 'dcb'), (SLC[:'m':-1], 'yx'),
-                (SLC['a':'a':-1], ''), (SLC['z':'z':-1], ''),
-                (SLC['m':'m':-1], '')
+            (SLC[::-1], 'yxdcb'), (SLC['b':'y':-1], ''),
+            (SLC['b'::-1], 'b'), (SLC[:'b':-1], 'yxdcb'),
+            (SLC[:'y':-1], 'y'), (SLC['y'::-1], 'yxdcb'),
+            (SLC['y'::-4], 'yb'),
+            # absent labels
+            (SLC[:'a':-1], 'yxdcb'), (SLC[:'a':-2], 'ydb'),
+            (SLC['z'::-1], 'yxdcb'), (SLC['z'::-3], 'yc'),
+            (SLC['m'::-1], 'dcb'), (SLC[:'m':-1], 'yx'),
+            (SLC['a':'a':-1], ''), (SLC['z':'z':-1], ''),
+            (SLC['m':'m':-1], '')
         ]:
             check_slice(in_slice, expected)
 

pandas/tests/indexes/test_category.py

@@ -207,19 +207,19 @@ def test_map(self):
         ci = pd.CategoricalIndex(list('ABABC'), categories=list('CBA'),
                                  ordered=True)
         result = ci.map(lambda x: x.lower())
-        exp = pd.Categorical(list('ababc'), categories=list('cba'),
-                             ordered=True)
-        tm.assert_categorical_equal(result, exp)
+        exp = pd.CategoricalIndex(list('ababc'), categories=list('cba'),
+                                  ordered=True)
+        tm.assert_index_equal(result, exp)
 
         ci = pd.CategoricalIndex(list('ABABC'), categories=list('BAC'),
                                  ordered=False, name='XXX')
         result = ci.map(lambda x: x.lower())
-        exp = pd.Categorical(list('ababc'), categories=list('bac'),
-                             ordered=False)
-        tm.assert_categorical_equal(result, exp)
+        exp = pd.CategoricalIndex(list('ababc'), categories=list('bac'),
+                                  ordered=False, name='XXX')
+        tm.assert_index_equal(result, exp)
 
-        tm.assert_numpy_array_equal(ci.map(lambda x: 1),
-                                    np.array([1] * 5, dtype=np.int64))
+        tm.assert_index_equal(ci.map(lambda x: 1),
+                              Index(np.array([1] * 5, dtype=np.int64), name='XXX'))
 
         # change categories dtype
         ci = pd.CategoricalIndex(list('ABABC'), categories=list('BAC'),
@@ -228,9 +228,9 @@ def f(x):
             return {'A': 10, 'B': 20, 'C': 30}.get(x)
 
         result = ci.map(f)
-        exp = pd.Categorical([10, 20, 10, 20, 30], categories=[20, 10, 30],
-                             ordered=False)
-        tm.assert_categorical_equal(result, exp)
+        exp = pd.CategoricalIndex([10, 20, 10, 20, 30], categories=[20, 10, 30],
+                                  ordered=False)
+        tm.assert_index_equal(result, exp)
 
     def test_where(self):
         i = self.create_index()

pandas/tests/test_categorical.py

@@ -1669,7 +1669,7 @@ def test_map(self):
         tm.assert_categorical_equal(result, exp)
 
         result = c.map(lambda x: 1)
-        tm.assert_numpy_array_equal(result, np.array([1] * 5, dtype=np.int64))
+        tm.assert_index_equal(result, Index(np.array([1] * 5, dtype=np.int64)))
 
 
 class TestCategoricalAsBlock(tm.TestCase):