From 7ff2b16fcc25e3a8031e92658355cfdd9fd55863 Mon Sep 17 00:00:00 2001 From: Andy Hayden Date: Sat, 2 Mar 2013 13:44:45 +0000 Subject: [PATCH 1/3] keep name after DataFrame drop --- pandas/core/generic.py | 9 +++++++-- pandas/tests/test_frame.py | 11 +++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 3ac3c8eef0a10..afe7f8775b1e9 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -340,7 +340,7 @@ def drop(self, labels, axis=0, level=None): dropped : type of caller """ axis_name = self._get_axis_name(axis) - axis = self._get_axis(axis) + axis, axis_ = self._get_axis(axis), axis if axis.is_unique: if level is not None: @@ -349,8 +349,13 @@ def drop(self, labels, axis=0, level=None): new_axis = axis.drop(labels, level=level) else: new_axis = axis.drop(labels) + dropped = self.reindex(**{axis_name: new_axis}) + try: + dropped.axes[axis_].names = axis.names + except AttributeError: + pass + return dropped - return self.reindex(**{axis_name: new_axis}) else: if level is not None: if not isinstance(axis, MultiIndex): diff --git a/pandas/tests/test_frame.py b/pandas/tests/test_frame.py index 9329bb1da2b07..a27bc581a2e69 100644 --- a/pandas/tests/test_frame.py +++ b/pandas/tests/test_frame.py @@ -5103,6 +5103,17 @@ def test_corrwith_series(self): assert_series_equal(result, expected) + def test_drop_names_axis(self): + df = DataFrame([[1, 2, 3],[3, 4, 5],[5, 6, 7]], index=['a', 'b', 'c'], columns=['d', 'e', 'f']) + dropped_b = DataFrame([[1, 2, 3],[5, 6, 7]], index=['a', 'c'], columns=['d', 'e', 'f']) + dropped_e = DataFrame([[1, 3], [3, 5], [5, 7]], index=['a', 'b', 'c'], columns=['d', 'f']) + for frame in [df, dropped_e, dropped_e]: + frame.index.name, frame.columns.name = 'first', 'second' + df_dropped_b = df.drop('b') + df_dropped_e = df.drop('e', axis=1) + assert_frame_equal(df_dropped_b, dropped_b) + assert_frame_equal(df_dropped_e, dropped_e) + def test_dropEmptyRows(self): N = len(self.frame.index) mat = randn(N) From 3f02a6a83e0609a246d371d5cdc08139bc0169ce Mon Sep 17 00:00:00 2001 From: Andy Hayden Date: Sun, 3 Mar 2013 00:40:22 +0000 Subject: [PATCH 2/3] fix testcase for drop losing names --- pandas/tests/test_frame.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/pandas/tests/test_frame.py b/pandas/tests/test_frame.py index a27bc581a2e69..69c8bdc60be13 100644 --- a/pandas/tests/test_frame.py +++ b/pandas/tests/test_frame.py @@ -5103,16 +5103,15 @@ def test_corrwith_series(self): assert_series_equal(result, expected) - def test_drop_names_axis(self): + def test_drop_names(self): df = DataFrame([[1, 2, 3],[3, 4, 5],[5, 6, 7]], index=['a', 'b', 'c'], columns=['d', 'e', 'f']) - dropped_b = DataFrame([[1, 2, 3],[5, 6, 7]], index=['a', 'c'], columns=['d', 'e', 'f']) - dropped_e = DataFrame([[1, 3], [3, 5], [5, 7]], index=['a', 'b', 'c'], columns=['d', 'f']) - for frame in [df, dropped_e, dropped_e]: - frame.index.name, frame.columns.name = 'first', 'second' + df.index.name, df.columns.name = 'first', 'second' df_dropped_b = df.drop('b') df_dropped_e = df.drop('e', axis=1) - assert_frame_equal(df_dropped_b, dropped_b) - assert_frame_equal(df_dropped_e, dropped_e) + self.assert_(df_dropped_b.index.name == 'first') + self.assert_(df_dropped_e.index.name == 'first') + self.assert_(df_dropped_b.columns.name == 'second') + self.assert_(df_dropped_e.columns.name == 'second') def test_dropEmptyRows(self): N = len(self.frame.index) From 64223074c0b7b3b6cd0d4225f3eab303bff6faea Mon Sep 17 00:00:00 2001 From: Andy Hayden Date: Sun, 3 Mar 2013 10:36:39 +0000 Subject: [PATCH 3/3] add name check to assert_frame_equal --- pandas/util/testing.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pandas/util/testing.py b/pandas/util/testing.py index 97059d9aaf9f5..6ce6a6d02c427 100644 --- a/pandas/util/testing.py +++ b/pandas/util/testing.py @@ -200,10 +200,12 @@ def assert_frame_equal(left, right, check_dtype=True, assert(type(left.index) == type(right.index)) assert(left.index.dtype == right.index.dtype) assert(left.index.inferred_type == right.index.inferred_type) + assert(left.index.names == right.index.names) if check_column_type: assert(type(left.columns) == type(right.columns)) assert(left.columns.dtype == right.columns.dtype) assert(left.columns.inferred_type == right.columns.inferred_type) + assert(left.columns.names == right.columns.names) def assert_panel_equal(left, right,