From 05de406d5e3c712e1bc0eddf77c65b794ed517fe Mon Sep 17 00:00:00 2001 From: Erfan Nariman Date: Mon, 11 Oct 2021 23:57:52 +0200 Subject: [PATCH 1/2] check if ndim == 1 and columns=tuple --- pandas/core/common.py | 12 ++++++++++-- pandas/core/indexes/base.py | 2 +- pandas/tests/frame/methods/test_drop.py | 7 +++++++ 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/pandas/core/common.py b/pandas/core/common.py index 2bf925466e176..21c66e9638cc8 100644 --- a/pandas/core/common.py +++ b/pandas/core/common.py @@ -248,7 +248,9 @@ def asarray_tuplesafe(values, dtype: NpDtype | None = None) -> np.ndarray: return result -def index_labels_to_array(labels, dtype: NpDtype | None = None) -> np.ndarray: +def index_labels_to_array( + labels, dtype: NpDtype | None = None, ndim: int | None = None +) -> np.ndarray: """ Transform label or iterable of labels to array, for use in Index. @@ -261,9 +263,15 @@ def index_labels_to_array(labels, dtype: NpDtype | None = None) -> np.ndarray: ------- array """ - if isinstance(labels, (str, tuple)): + if isinstance(labels, str): labels = [labels] + if isinstance(labels, tuple): + if ndim == 1: + labels = list(labels) + else: + labels = [labels] + if not isinstance(labels, (list, np.ndarray)): try: labels = list(labels) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index a6360454cc02c..7868b2ca3b3ce 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -6384,7 +6384,7 @@ def drop(self, labels, errors: str_t = "raise") -> Index: if not isinstance(labels, Index): # avoid materializing e.g. RangeIndex arr_dtype = "object" if self.dtype == "object" else None - labels = com.index_labels_to_array(labels, dtype=arr_dtype) + labels = com.index_labels_to_array(labels, dtype=arr_dtype, ndim=self.ndim) indexer = self.get_indexer_for(labels) mask = indexer == -1 diff --git a/pandas/tests/frame/methods/test_drop.py b/pandas/tests/frame/methods/test_drop.py index 1e2ce38a2fefd..df6b7bca0b92e 100644 --- a/pandas/tests/frame/methods/test_drop.py +++ b/pandas/tests/frame/methods/test_drop.py @@ -546,3 +546,10 @@ def test_drop_level_missing_label_multiindex(self): df = DataFrame(index=MultiIndex.from_product([range(3), range(3)])) with pytest.raises(KeyError, match="labels \\[5\\] not found in level"): df.drop(5, level=0) + + def test_drop_labels_is_tuple(self): + # GH 43978 + df = DataFrame(np.arange(12).reshape(3, 4), columns=["A", "B", "C", "D"]) + result = df.drop(columns=("A", "B")) + expected = DataFrame({"C": [2, 6, 10], "D": [3, 7, 11]}) + tm.assert_frame_equal(result, expected) From d1fea9d3170c0897b92f7a5c9472cf7583c2180f Mon Sep 17 00:00:00 2001 From: Erfan Nariman Date: Tue, 12 Oct 2021 00:12:51 +0200 Subject: [PATCH 2/2] add test drop with actual tuple columns --- pandas/tests/frame/methods/test_drop.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pandas/tests/frame/methods/test_drop.py b/pandas/tests/frame/methods/test_drop.py index df6b7bca0b92e..9050e165e9365 100644 --- a/pandas/tests/frame/methods/test_drop.py +++ b/pandas/tests/frame/methods/test_drop.py @@ -553,3 +553,10 @@ def test_drop_labels_is_tuple(self): result = df.drop(columns=("A", "B")) expected = DataFrame({"C": [2, 6, 10], "D": [3, 7, 11]}) tm.assert_frame_equal(result, expected) + + def test_drop_columns_are_tuple(self): + # GH 43978 + df = DataFrame(np.arange(6).reshape(3, 2), columns=[("A", "B"), ("C", "D")]) + result = df.drop(columns=("A", "B")) + expected = DataFrame([1, 3, 5], columns=[("C", "D")]) + tm.assert_frame_equal(result, expected)