Merge pull request #5354 from jtratner/allow-drop-to-accept-names

BUG/ENH: Allow drop to accept axis name

Author: jtratner

doc/source/release.rst

@@ -189,6 +189,8 @@ Improvements to existing features
     right-open fashion using the ``closed`` parameter (:issue:`4579`)
   - Python csv parser now supports usecols (:issue:`4335`)
   - Added support for Google Analytics v3 API segment IDs that also supports v2 IDs. (:issue:`5271`)
+  - ``NDFrame.drop()`` now accepts names as well as integers for the axis
+    argument. (:issue:`5354`)
 
 API Changes
 ~~~~~~~~~~~

pandas/core/generic.py

@@ -1093,14 +1093,15 @@ def drop(self, labels, axis=0, level=None):
         Parameters
         ----------
         labels : single label or list-like
-        axis : int
+        axis : int or axis name
         level : int or name, default None
             For MultiIndex
 
         Returns
         -------
         dropped : type of caller
         """
+        axis = self._get_axis_number(axis)
         axis_name = self._get_axis_name(axis)
         axis, axis_ = self._get_axis(axis), axis
 

pandas/tests/test_frame.py

@@ -6734,6 +6734,25 @@ def test_drop_col_still_multiindex(self):
         del df[('a', '', '')]
         assert(isinstance(df.columns, MultiIndex))
 
+    def test_drop(self):
+        simple = DataFrame({"A": [1, 2, 3, 4], "B": [0, 1, 2, 3]})
+        assert_frame_equal(simple.drop("A", axis=1), simple[['B']])
+        assert_frame_equal(simple.drop(["A", "B"], axis='columns'),
+                           simple[[]])
+        assert_frame_equal(simple.drop([0, 1, 3], axis=0), simple.ix[[2], :])
+        assert_frame_equal(simple.drop([0, 3], axis='index'), simple.ix[[1, 2], :])
+
+        #non-unique - wheee!
+        nu_df = DataFrame(lzip(range(3), range(-3, 1), list('abc')),
+                          columns=['a', 'a', 'b'])
+        assert_frame_equal(nu_df.drop('a', axis=1), nu_df[['b']])
+        assert_frame_equal(nu_df.drop('b', axis='columns'), nu_df['a'])
+
+        nu_df = nu_df.set_index(pd.Index(['X', 'Y', 'X']))
+        nu_df.columns = list('abc')
+        assert_frame_equal(nu_df.drop('X', axis='rows'), nu_df.ix[["Y"], :])
+        assert_frame_equal(nu_df.drop(['X', 'Y'], axis=0), nu_df.ix[[], :])
+
     def test_fillna(self):
         self.tsframe['A'][:5] = nan
         self.tsframe['A'][-5:] = nan

pandas/tests/test_panel.py

@@ -1624,6 +1624,43 @@ def test_dropna(self):
         exp = p.ix[['a', 'c', 'd']]
         assert_panel_equal(result, exp)
 
+    def test_drop(self):
+        df = DataFrame({"A": [1, 2], "B": [3, 4]})
+        panel = Panel({"One": df, "Two": df})
+
+        def check_drop(drop_val, axis_number, aliases, expected):
+            try:
+                actual = panel.drop(drop_val, axis=axis_number)
+                assert_panel_equal(actual, expected)
+                for alias in aliases:
+                    actual = panel.drop(drop_val, axis=alias)
+                    assert_panel_equal(actual, expected)
+            except AssertionError:
+                print("Failed with axis_number %d and aliases: %s" %
+                      (axis_number, aliases))
+                raise
+        # Items
+        expected = Panel({"One": df})
+        check_drop('Two', 0, ['items'], expected)
+
+        # Major
+        exp_df = DataFrame({"A": [2], "B": [4]}, index=[1])
+        expected = Panel({"One": exp_df, "Two": exp_df})
+        check_drop(0, 1, ['major_axis', 'major'], expected)
+
+        exp_df = DataFrame({"A": [1], "B": [3]}, index=[0])
+        expected = Panel({"One": exp_df, "Two": exp_df})
+        check_drop([1], 1, ['major_axis', 'major'], expected)
+
+        # Minor
+        exp_df = df[['B']]
+        expected = Panel({"One": exp_df, "Two": exp_df})
+        check_drop(["A"], 2, ['minor_axis', 'minor'], expected)
+
+        exp_df = df[['A']]
+        expected = Panel({"One": exp_df, "Two": exp_df})
+        check_drop("B", 2, ['minor_axis', 'minor'], expected)
+
     def test_update(self):
         pan = Panel([[[1.5, np.nan, 3.],
                     [1.5, np.nan, 3.],

pandas/tests/test_series.py

@@ -1429,14 +1429,14 @@ def test_drop(self):
         expected = Series([1],index=['one'])
         result = s.drop(['two'])
         assert_series_equal(result,expected)
-        result = s.drop('two')
+        result = s.drop('two', axis='rows')
         assert_series_equal(result,expected)
 
         # non-unique
         # GH 5248
         s = Series([1,1,2],index=['one','two','one'])
         expected = Series([1,2],index=['one','one'])
-        result = s.drop(['two'])
+        result = s.drop(['two'], axis=0)
         assert_series_equal(result,expected)
         result = s.drop('two')
         assert_series_equal(result,expected)
@@ -1452,6 +1452,9 @@ def test_drop(self):
         self.assertRaises(ValueError, s.drop, 'bc')
         self.assertRaises(ValueError, s.drop, ('a',))
 
+        # bad axis
+        self.assertRaises(ValueError, s.drop, 'one', axis='columns')
+
     def test_ix_setitem(self):
         inds = self.series.index[[3, 4, 7]]