ENH/API: DataFrame.stack() support for level=ALL_LEVELS and sequentially=True/False.

Author: seth-p

pandas/core/frame.py

@@ -3392,7 +3392,7 @@ def pivot(self, index=None, columns=None, values=None):
         from pandas.core.reshape import pivot
         return pivot(self, index=index, columns=columns, values=values)
 
-    def stack(self, level=-1, dropna=True):
+    def stack(self, level=-1, dropna=True, sequentially=True):
         """
         Pivot a level of the (possibly hierarchical) column labels, returning a
         DataFrame (or Series in the case of an object with a single level of
@@ -3403,10 +3403,14 @@ def stack(self, level=-1, dropna=True):
         Parameters
         ----------
         level : int, string, or list of these, default last level
-            Level(s) to stack, can pass level name
+            Level(s) to stack, can pass level name(s).
+            May pass Index.ALL_LEVELS to specify list(range(columns.nlevels)).
         dropna : boolean, default True
             Whether to drop rows in the resulting Frame/Series with no valid
             values
+        sequentially : boolean, default True
+            When level is a list (or ALL_LEVELS), whether the multiple column levels
+            should be stacked sequentially (if True) or simultaneously (if False).
 
         Examples
         ----------
@@ -3425,12 +3429,13 @@ def stack(self, level=-1, dropna=True):
         -------
         stacked : DataFrame or Series
         """
-        from pandas.core.reshape import stack, stack_multiple
+        from pandas.core.reshape import stack_levels_sequentially, stack_multi_levels_simultaneously
 
-        if isinstance(level, (tuple, list)):
-            return stack_multiple(self, level, dropna=dropna)
+        level_num = self.columns._get_level_numbers(level, allow_mixed_names_and_numbers=False)
+        if isinstance(level_num, (tuple, list, set)) and isinstance(self.columns, MultiIndex) and (not sequentially):
+                return stack_multi_levels_simultaneously(self, level_num, dropna=dropna)
         else:
-            return stack(self, level, dropna=dropna)
+            return stack_levels_sequentially(self, level_num, dropna=dropna)
 
     def unstack(self, level=-1):
         """

pandas/core/index.py

@@ -123,6 +123,8 @@ class Index(IndexOpsMixin, PandasObject):
 
     _engine_type = _index.ObjectEngine
 
+    ALL_LEVELS = -1000
+
     def __new__(cls, data=None, dtype=None, copy=False, name=None, fastpath=False,
                 tupleize_cols=True, **kwargs):
 
@@ -824,10 +826,26 @@ def _validate_index_level(self, level):
             raise KeyError('Level %s must be same as name (%s)'
                            % (level, self.name))
 
-    def _get_level_number(self, level):
+    def _get_level_number(self, level, ignore_names=False):
+        if ignore_names and (not isinstance(level, int)):
+            raise KeyError('Level %s not found' % str(level))
         self._validate_index_level(level)
         return 0
 
+    def _get_level_numbers(self, levels, allow_mixed_names_and_numbers=False):
+        if levels == self.ALL_LEVELS:
+            return list(range(self.nlevels))
+        elif isinstance(levels, (list, tuple, set)):
+            if (not allow_mixed_names_and_numbers) and (not all(lev in self.names for lev in levels)):
+                if all(isinstance(lev, int) for lev in levels):
+                    return type(levels)(self._get_level_number(level, ignore_names=True) for level in levels)
+                else:
+                    raise ValueError("level should contain all level names or all level numbers, "
+                                     "not a mixture of the two.")
+            return type(levels)(self._get_level_number(level) for level in levels)
+        else:
+            return self._get_level_number(levels)
+
     @cache_readonly
     def inferred_type(self):
         """ return a string of the type inferred from the values """
@@ -3161,28 +3179,30 @@ def _from_elements(values, labels=None, levels=None, names=None,
                        sortorder=None):
         return MultiIndex(levels, labels, names, sortorder=sortorder)
 
-    def _get_level_number(self, level):
-        try:
+    def _get_level_number(self, level, ignore_names=False):
+        if not ignore_names:
             count = self.names.count(level)
             if count > 1:
                 raise ValueError('The name %s occurs multiple times, use a '
                                  'level number' % level)
-            level = self.names.index(level)
-        except ValueError:
-            if not isinstance(level, int):
-                raise KeyError('Level %s not found' % str(level))
-            elif level < 0:
-                level += self.nlevels
-                if level < 0:
-                    orig_level = level - self.nlevels
-                    raise IndexError(
-                        'Too many levels: Index has only %d levels, '
-                        '%d is not a valid level number' % (self.nlevels, orig_level)
-                    )
-            # Note: levels are zero-based
-            elif level >= self.nlevels:
-                raise IndexError('Too many levels: Index has only %d levels, '
-                                 'not %d' % (self.nlevels, level + 1))
+            try:
+                return self.names.index(level)
+            except ValueError:
+                pass
+        if not isinstance(level, int):
+            raise KeyError('Level %s not found' % str(level))
+        elif level < 0:
+            level += self.nlevels
+            if level < 0:
+                orig_level = level - self.nlevels
+                raise IndexError(
+                    'Too many levels: Index has only %d levels, '
+                    '%d is not a valid level number' % (self.nlevels, orig_level)
+                )
+        # Note: levels are zero-based
+        elif level >= self.nlevels:
+            raise IndexError('Too many levels: Index has only %d levels, '
+                             'not %d' % (self.nlevels, level + 1))
         return level
 
     _tuples = None
@@ -4852,7 +4872,7 @@ def _trim_front(strings):
 
 
 def _sanitize_and_check(indexes):
-    kinds = list(set([type(index) for index in indexes]))
+    kinds = list(set(type(index) for index in indexes))
 
     if list in kinds:
         if len(kinds) > 1:
@@ -4873,11 +4893,11 @@ def _get_consensus_names(indexes):
 
     # find the non-none names, need to tupleify to make
     # the set hashable, then reverse on return
-    consensus_names = set([
+    consensus_names = set(
         tuple(i.names) for i in indexes if all(n is not None for n in i.names)
-    ])
+    )
     if len(consensus_names) == 1:
-        return list(list(consensus_names)[0])
+        return list(consensus_names.pop())
     return [None] * indexes[0].nlevels