From aba8bd763141b6bafb5fd600bdedacf6ddbbbe51 Mon Sep 17 00:00:00 2001 From: Jeff Tratner Date: Sun, 27 Oct 2013 20:09:04 -0400 Subject: [PATCH] BUG/ENH: Allow drop to accept axis name Add tests for DataFrame, Series and Panel, with axis checks, etc. --- doc/source/release.rst | 2 ++ pandas/core/generic.py | 3 ++- pandas/tests/test_frame.py | 19 +++++++++++++++++++ pandas/tests/test_panel.py | 37 +++++++++++++++++++++++++++++++++++++ pandas/tests/test_series.py | 7 +++++-- 5 files changed, 65 insertions(+), 3 deletions(-) diff --git a/doc/source/release.rst b/doc/source/release.rst index c544396e4c73b..5ca602cc1cc0f 100644 --- a/doc/source/release.rst +++ b/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 ~~~~~~~~~~~ diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 3259522e20926..704c41819ae71 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -1093,7 +1093,7 @@ 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 @@ -1101,6 +1101,7 @@ def drop(self, labels, axis=0, level=None): ------- dropped : type of caller """ + axis = self._get_axis_number(axis) axis_name = self._get_axis_name(axis) axis, axis_ = self._get_axis(axis), axis diff --git a/pandas/tests/test_frame.py b/pandas/tests/test_frame.py index 2d214a2abe867..6c7877a8d6e4a 100644 --- a/pandas/tests/test_frame.py +++ b/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 diff --git a/pandas/tests/test_panel.py b/pandas/tests/test_panel.py index d2da403907280..39ceba7469f36 100644 --- a/pandas/tests/test_panel.py +++ b/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.], diff --git a/pandas/tests/test_series.py b/pandas/tests/test_series.py index 6abfce31c4870..ad6fa68f063e1 100644 --- a/pandas/tests/test_series.py +++ b/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]]