Merge pull request #8840 from jorisvandenbossche/gh5620

TST: add tests for index.groupby (GH5620)

Author: jorisvandenbossche

pandas/core/index.py

@@ -1555,6 +1555,20 @@ def _possibly_promote(self, other):
         return self, other
 
     def groupby(self, to_groupby):
+        """
+        Group the index labels by a given array of values.
+
+        Parameters
+        ----------
+        to_groupby : array
+            Values used to determine the groups.
+
+        Returns
+        -------
+        groups : dict
+            {group name -> group labels}
+
+        """
         return self._groupby(self.values, _values_from_object(to_groupby))
 
     def map(self, mapper):

pandas/tests/test_index.py

@@ -1137,6 +1137,11 @@ def get_reindex_type(target):
         self.assertEqual(reindexed.levels[0].dtype.type, np.int64)
         self.assertEqual(reindexed.levels[1].dtype.type, np.float64)
 
+    def test_groupby(self):
+        idx = Index(range(5))
+        groups = idx.groupby(np.array([1,1,2,2,2]))
+        exp = {1: [0, 1], 2: [2, 3, 4]}
+        tm.assert_dict_equal(groups, exp)
 
 
 class Numeric(Base):
@@ -3524,6 +3529,17 @@ def test_reindex_lvl_preserves_type_if_target_is_empty_list_or_array(self):
         self.assertEqual(idx.reindex([], level=1)[0].levels[1].dtype.type,
                          np.object_)
 
+    def test_groupby(self):
+        groups = self.index.groupby(np.array([1, 1, 1, 2, 2, 2]))
+        labels = self.index.get_values().tolist()
+        exp = {1: labels[:3], 2: labels[3:]}
+        tm.assert_dict_equal(groups, exp)
+
+        # GH5620
+        groups = self.index.groupby(self.index)
+        exp = dict((key, [key]) for key in self.index)
+        tm.assert_dict_equal(groups, exp)
+
 
 def test_get_combined_index():
     from pandas.core.index import _get_combined_index