Cleaned up _get_grouper_for_level implementations and added docstrings

Author: Jon M. Mease

pandas/indexes/base.py

@@ -432,15 +432,33 @@ def _update_inplace(self, result, **kwargs):
         # guard when called from IndexOpsMixin
         raise TypeError("Index can't be updated inplace")
 
-    def _get_grouper_for_level(self, grouper, level):
-        # Use self (Index) as grouper if None was passed
-        if grouper is None:
+    def _get_grouper_for_level(self, group_mapper, level):
+        """
+        Get index grouper corresponding to an index level
+
+        Parameters
+        ----------
+        group_mapper: Group mapping function or None
+            Function mapping index values to groups
+        level : int
+            Index level (Only used by MultiIndex override)
+
+        Returns
+        -------
+        grouper : Index
+            Index of values to group on
+        labels : None
+            Array of locations in level_index
+            (Only returned by MultiIndex override)
+        level_index : None
+            Index of unique values for level
+            (Only returned by MultiIndex override)
+        """
+        if group_mapper is None:
             grouper = self
+        else:
+            grouper = self.map(group_mapper)
 
-        # Return tuple of (grouper, labels, level_index)
-        # where labels and level_index are None for the Index
-        # implementation.  The labels and level_index values
-        # are only calculated in the MultiIndex implementation
         return grouper, None, None
 
     def is_(self, other):

pandas/indexes/multi.py

@@ -524,36 +524,52 @@ def _format_native_types(self, na_rep='nan', **kwargs):
 
         return mi.values
 
-    def _get_grouper_for_level(self, grouper, level):
+    def _get_grouper_for_level(self, group_mapper, level):
+        """
+        Get index grouper corresponding to an index level
 
+        Parameters
+        ----------
+        group_mapper: Group mapping function or None
+            Function mapping index values to groups
+        level : int
+            Index level
+
+        Returns
+        -------
+        grouper : Index
+            Index of values to group on
+        labels : ndarray of int or None
+            Array of locations in level_index
+        level_index : Index or None
+            Index of unique values for level
+        """
         inds = self.labels[level]
         level_index = self.levels[level]
 
-        # XXX complete hack
-
-        if grouper is not None:
+        if group_mapper is not None:
+            # Handle group mapping function and return
             level_values = self.levels[level].take(inds)
-            grouper = level_values.map(grouper)
-            labels = None
-            level_index = None
-        else:
-            # all levels may not be observed
-            labels, uniques = algos.factorize(inds, sort=True)
+            grouper = level_values.map(group_mapper)
+            return grouper, None, None
+
+        labels, uniques = algos.factorize(inds, sort=True)
 
-            if len(uniques) > 0 and uniques[0] == -1:
-                # handle NAs
-                mask = inds != -1
-                ok_labels, uniques = algos.factorize(inds[mask],
-                                                     sort=True)
+        if len(uniques) > 0 and uniques[0] == -1:
+            # Handle NAs
+            mask = inds != -1
+            ok_labels, uniques = algos.factorize(inds[mask],
+                                                 sort=True)
 
-                labels = np.empty(len(inds), dtype=inds.dtype)
-                labels[mask] = ok_labels
-                labels[~mask] = -1
+            labels = np.empty(len(inds), dtype=inds.dtype)
+            labels[mask] = ok_labels
+            labels[~mask] = -1
 
-            if len(uniques) < len(level_index):
-                level_index = level_index.take(uniques)
+        if len(uniques) < len(level_index):
+            # Remove unobserved levels from level_index
+            level_index = level_index.take(uniques)
 
-            grouper = level_index.take(labels)
+        grouper = level_index.take(labels)
 
         return grouper, labels, level_index