ENH: xs level can take multiple levels, pass multiple levels to MultiIndex.droplevel, GH #371

Author: wesm

pandas/core/frame.py

@@ -1426,6 +1426,7 @@ def xs(self, key, axis=0, level=None, copy=True):
                 indexer = loc
 
             result = self.ix[indexer]
+
             new_ax = result._get_axis(axis).droplevel(level)
             setattr(result, result._get_axis_name(axis), new_ax)
             return result

pandas/core/index.py

@@ -1394,7 +1394,7 @@ def droplevel(self, level=0):
 
         Parameters
         ----------
-        level : int
+        level : int/level name or list thereof
 
         Notes
         -----
@@ -1404,12 +1404,20 @@ def droplevel(self, level=0):
         -------
         index : Index or MultiIndex
         """
+        levels = level
+        if not isinstance(levels, (tuple, list)):
+            levels = [level]
+
         new_levels = list(self.levels)
-        new_levels.pop(level)
         new_labels = list(self.labels)
-        new_labels.pop(level)
         new_names = list(self.names)
-        new_names.pop(level)
+
+        levnums = sorted(self._get_level_number(lev) for lev in levels)[::-1]
+
+        for i in levnums:
+            new_levels.pop(i)
+            new_labels.pop(i)
+            new_names.pop(i)
 
         if len(new_levels) == 1:
             result = new_levels[0].take(new_labels[0])
@@ -1692,6 +1700,21 @@ def get_loc_level(self, key, level=0):
         -------
         loc : int or slice object
         """
+        if isinstance(level, (tuple, list)):
+            assert(len(key) == len(level))
+            result = None
+            for lev, k in zip(level, key):
+                loc = self.get_loc_level(k, level=lev)
+                if isinstance(loc, slice):
+                    mask = np.zeros(len(self), dtype=bool)
+                    mask[loc] = True
+                    loc = mask
+
+                result = loc if result is None else result & loc
+            return result
+
+        level = self._get_level_number(level)
+
         if isinstance(key, tuple) and level == 0:
             if not any(isinstance(k, slice) for k in key):
                 if len(key) == self.nlevels:

pandas/tests/test_index.py

@@ -1208,6 +1208,23 @@ def test_droplevel_with_names(self):
         dropped = index.droplevel(0)
         self.assertEqual(dropped.names, ['two', 'three'])
 
+        dropped = index.droplevel('two')
+        expected = index.droplevel(1)
+        self.assert_(dropped.equals(expected))
+
+    def test_droplevel_multiple(self):
+        index = MultiIndex(levels=[Index(range(4)),
+                                   Index(range(4)),
+                                   Index(range(4))],
+                           labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]),
+                                   np.array([0, 1, 0, 0, 0, 1, 0, 1]),
+                                   np.array([1, 0, 1, 1, 0, 0, 1, 0])],
+                           names=['one', 'two', 'three'])
+
+        dropped = index[:2].droplevel(['three', 'one'])
+        expected = index[:2].droplevel(2).droplevel(0)
+        self.assert_(dropped.equals(expected))
+
     def test_insert(self):
         # key contained in all levels
         new_index = self.index.insert(0, ('bar', 'two'))

pandas/tests/test_multilevel.py

@@ -229,7 +229,7 @@ def test_xs_partial(self):
         assert_frame_equal(result, result2)
 
     def test_xs_level(self):
-        result = self.frame.xs('two', level=1)
+        result = self.frame.xs('two', level='second')
         expected = self.frame[self.frame.index.get_level_values(1) == 'two']
         expected.index = expected.index.droplevel(1)
 
@@ -243,6 +243,21 @@ def test_xs_level(self):
         expected.index = expected.index.droplevel(2)
         assert_frame_equal(result, expected)
 
+    def test_xs_level_multiple(self):
+        from pandas import read_table
+        from StringIO import StringIO
+        text = """                      A       B       C       D        E
+one two three   four
+a   b   10.0032 5    -0.5109 -2.3358 -0.4645  0.05076  0.3640
+a   q   20      4     0.4473  1.4152  0.2834  1.00661  0.1744
+x   q   30      3    -0.6662 -0.5243 -0.3580  0.89145  2.5838"""
+
+        df = read_table(StringIO(text), sep='\s+')
+
+        result = df.xs(('a', 4), level=['one', 'four'])
+        expected = df.xs('a').xs(4, level='four')
+        assert_frame_equal(result, expected)
+
     def test_xs_level0(self):
         from pandas import read_table
         from StringIO import StringIO