ENH: unstack/stack multiple levels per #370, use Series index per note in #373

Author: wesm

pandas/core/frame.py

@@ -15,7 +15,6 @@
 from itertools import izip
 from StringIO import StringIO
 import csv
-import gc
 import operator
 import sys
 
@@ -241,6 +240,10 @@ def _init_ndarray(self, values, index, columns, dtype=None,
         if isinstance(values, Series) and values.name is not None:
             if columns is None:
                 columns = [values.name]
+            if index is None:
+                index = values.index
+            else:
+                values = values.reindex(index)
 
         values = _prep_ndarray(values, copy=copy)
 
@@ -1958,24 +1961,31 @@ def stack(self, level=-1, dropna=True):
 
         Parameters
         ----------
-        level : int or string, default last level
-            Level to stack, can pass level name
+        level : int, string, or list of these, default last level
+            Level(s) to stack, can pass level name
 
         Returns
         -------
         stacked : Series
         """
         from pandas.core.reshape import stack
-        return stack(self, level=level, dropna=dropna)
+
+        if isinstance(level, (tuple, list)):
+            result = self
+            for lev in level:
+                result = stack(result, lev, dropna=dropna)
+            return result
+        else:
+            return stack(self, level, dropna=dropna)
 
     def unstack(self, level=-1):
         """
         "Unstack" level from MultiLevel index to produce reshaped DataFrame
 
         Parameters
         ----------
-        level : int or string, default last level
-            Level to unstack, can pass level name
+        level : int, string, or list of these, default last level
+            Level(s) to unstack, can pass level name
 
         Examples
         --------
@@ -1999,10 +2009,14 @@ def unstack(self, level=-1):
         -------
         unstacked : DataFrame
         """
-        from pandas.core.reshape import _Unstacker
-        unstacker = _Unstacker(self.values, self.index, level=level,
-                               value_columns=self.columns)
-        return unstacker.get_result()
+        from pandas.core.reshape import unstack
+        if isinstance(level, (tuple, list)):
+            result = self
+            for lev in level:
+                result = unstack(result, lev)
+            return result
+        else:
+            return unstack(self, level)
 
     def delevel(self):
         """

pandas/core/reshape.py

@@ -11,6 +11,7 @@
 from pandas.core.common import notnull
 from pandas.core.index import MultiIndex
 
+
 class ReshapeError(Exception):
     pass
 
@@ -280,6 +281,15 @@ def _slow_pivot(index, columns, values):
 
     return DataFrame(tree)
 
+def unstack(obj, level):
+    if isinstance(obj, DataFrame):
+        columns = obj.columns
+    else:
+        columns = None
+    unstacker = _Unstacker(obj.values, obj.index, level=level,
+                           value_columns=columns)
+    return unstacker.get_result()
+
 def stack(frame, level=-1, dropna=True):
     """
     Convert DataFrame to Series with multi-level Index. Columns become the

pandas/core/series.py

@@ -1331,8 +1331,8 @@ def unstack(self, level=-1):
 
         Parameters
         ----------
-        level : int, default last level
-            Level to unstack
+        level : int, string, or list of these, default last level
+            Level(s) to unstack, can pass level name
 
         Examples
         --------
@@ -1356,9 +1356,14 @@ def unstack(self, level=-1):
         -------
         unstacked : DataFrame
         """
-        from pandas.core.reshape import _Unstacker
-        unstacker = _Unstacker(self.values, self.index, level=level)
-        return unstacker.get_result()
+        from pandas.core.reshape import unstack
+        if isinstance(level, (tuple, list)):
+            result = self
+            for lev in level:
+                result = unstack(result, lev)
+            return result
+        else:
+            return unstack(self, level)
 
     #----------------------------------------------------------------------
     # function application

pandas/tests/test_frame.py

@@ -1147,6 +1147,7 @@ def test_constructor_Series_named(self):
         a = Series([1,2,3], index=['a','b','c'], name='x')
         df = DataFrame(a)
         self.assert_(df.columns[0] == 'x')
+        self.assert_(df.index.equals(a.index))
 
     def test_astype(self):
         casted = self.frame.astype(int)

pandas/tests/test_multilevel.py

@@ -55,6 +55,7 @@ def setUp(self):
         self.tdf = tm.makeTimeDataFrame()
         self.ymd = self.tdf.groupby([lambda x: x.year, lambda x: x.month,
                                      lambda x: x.day]).sum()
+        self.ymd.index.names = ['year', 'month', 'day']
 
     def test_append(self):
         a, b = self.frame[:5], self.frame[5:]
@@ -450,6 +451,26 @@ def test_stack_level_name(self):
         expected = self.frame.stack()
         assert_series_equal(result, expected)
 
+    def test_stack_unstack_multiple(self):
+        unstacked = self.ymd.unstack(['year', 'month'])
+        expected = self.ymd.unstack('year').unstack('month')
+        assert_frame_equal(unstacked, expected)
+        self.assertEquals(unstacked.columns.names,
+                          expected.columns.names)
+
+        # series
+        s = self.ymd['A']
+        s_unstacked = s.unstack(['year', 'month'])
+        assert_frame_equal(s_unstacked, expected['A'])
+
+        restacked = unstacked.stack(['year', 'month'])
+        restacked = restacked.swaplevel(0, 1).swaplevel(1, 2)
+        restacked = restacked.sortlevel(0)
+
+        assert_frame_equal(restacked, self.ymd)
+        self.assertEquals(restacked.index.names,
+                          self.ymd.index.names)
+
     def test_groupby_transform(self):
         s = self.frame['A']
         grouper = s.index.get_level_values(0)