BUG: support for "level=" when reset_index() is called with a flat Index

closes #16263

Author: toobaz

doc/source/whatsnew/v0.20.2.txt

@@ -42,7 +42,7 @@ Conversion
 Indexing
 ^^^^^^^^
 
-
+- Bug in ``DataFrame.reset_index(level=)`` with flat index (:issue:`16263`)
 
 
 I/O

pandas/core/frame.py

@@ -3012,12 +3012,12 @@ def _maybe_casted_values(index, labels=None):
             return values
 
         new_index = _default_index(len(new_obj))
-        if isinstance(self.index, MultiIndex):
-            if level is not None:
-                if not isinstance(level, (tuple, list)):
-                    level = [level]
-                level = [self.index._get_level_number(lev) for lev in level]
-                if len(level) < len(self.index.levels):
+        if level is not None:
+            if not isinstance(level, (tuple, list)):
+                level = [level]
+            level = [self.index._get_level_number(lev) for lev in level]
+            if isinstance(self.index, MultiIndex):
+                if len(level) < self.index.nlevels:
                     new_index = self.index.droplevel(level)
 
         if not drop:
@@ -3033,6 +3033,8 @@ def _maybe_casted_values(index, labels=None):
 
             multi_col = isinstance(self.columns, MultiIndex)
             for i, (lev, lab) in reversed(list(enumerate(to_insert))):
+                if not (level is None or i in level):
+                    continue
                 name = names[i]
                 if multi_col:
                     col_name = (list(name) if isinstance(name, tuple)
@@ -3049,11 +3051,9 @@ def _maybe_casted_values(index, labels=None):
                     missing = self.columns.nlevels - len(name_lst)
                     name_lst += [col_fill] * missing
                     name = tuple(name_lst)
-
                 # to ndarray and maybe infer different dtype
                 level_values = _maybe_casted_values(lev, lab)
-                if level is None or i in level:
-                    new_obj.insert(0, name, level_values)
+                new_obj.insert(0, name, level_values)
 
         new_obj.index = new_index
         if not inplace:

pandas/tests/frame/test_alter_axes.py

@@ -641,6 +641,36 @@ def test_reset_index(self):
         xp = xp.set_index(['B'], append=True)
         assert_frame_equal(rs, xp, check_names=False)
 
+    def test_reset_index_level(self):
+        df = pd.DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]],
+                          columns=['A', 'B', 'C', 'D'])
+
+        # With MultiIndex
+        for levels in ['A', 'B'], [0, 1]:
+            result = df.set_index(['A', 'B']).reset_index(level=levels[0])
+            tm.assert_frame_equal(result, df.set_index('B'))
+
+            result = df.set_index(['A', 'B']).reset_index(level=levels[:1])
+            tm.assert_frame_equal(result, df.set_index('B'))
+
+            result = df.set_index(['A', 'B']).reset_index(level=levels)
+            tm.assert_frame_equal(result, df)
+
+        # With flat Index (GH 16263)
+        for level in 'A', 0:
+            result = df.set_index('A').reset_index(level=level)
+            tm.assert_frame_equal(result, df)
+
+            result = df.set_index('A').reset_index(level=[level])
+            tm.assert_frame_equal(result, df)
+
+        # Missing levels - for both MultiIndex and flat Index:
+        for idx_lev in ['A', 'B'], ['A']:
+            with tm.assert_raises_regex(KeyError, 'Level E '):
+                df.set_index(idx_lev).reset_index(level=['A', 'E'])
+            with tm.assert_raises_regex(IndexError, 'Too many levels'):
+                df.set_index(idx_lev).reset_index(level=[0, 1, 2])
+
     def test_reset_index_right_dtype(self):
         time = np.arange(0.0, 10, np.sqrt(2) / 2)
         s1 = Series((9.81 * time ** 2) / 2,