61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
import re
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from pandas import (
|
|
DataFrame,
|
|
MultiIndex,
|
|
)
|
|
|
|
|
|
class TestDataFrameDelItem:
|
|
def test_delitem(self, float_frame):
|
|
del float_frame["A"]
|
|
assert "A" not in float_frame
|
|
|
|
def test_delitem_multiindex(self):
|
|
midx = MultiIndex.from_product([["A", "B"], [1, 2]])
|
|
df = DataFrame(np.random.randn(4, 4), columns=midx)
|
|
assert len(df.columns) == 4
|
|
assert ("A",) in df.columns
|
|
assert "A" in df.columns
|
|
|
|
result = df["A"]
|
|
assert isinstance(result, DataFrame)
|
|
del df["A"]
|
|
|
|
assert len(df.columns) == 2
|
|
|
|
# A still in the levels, BUT get a KeyError if trying
|
|
# to delete
|
|
assert ("A",) not in df.columns
|
|
with pytest.raises(KeyError, match=re.escape("('A',)")):
|
|
del df[("A",)]
|
|
|
|
# behavior of dropped/deleted MultiIndex levels changed from
|
|
# GH 2770 to GH 19027: MultiIndex no longer '.__contains__'
|
|
# levels which are dropped/deleted
|
|
assert "A" not in df.columns
|
|
with pytest.raises(KeyError, match=re.escape("('A',)")):
|
|
del df["A"]
|
|
|
|
def test_delitem_corner(self, float_frame):
|
|
f = float_frame.copy()
|
|
del f["D"]
|
|
assert len(f.columns) == 3
|
|
with pytest.raises(KeyError, match=r"^'D'$"):
|
|
del f["D"]
|
|
del f["B"]
|
|
assert len(f.columns) == 2
|
|
|
|
def test_delitem_col_still_multiindex(self):
|
|
arrays = [["a", "b", "c", "top"], ["", "", "", "OD"], ["", "", "", "wx"]]
|
|
|
|
tuples = sorted(zip(*arrays))
|
|
index = MultiIndex.from_tuples(tuples)
|
|
|
|
df = DataFrame(np.random.randn(3, 4), columns=index)
|
|
del df[("a", "", "")]
|
|
assert isinstance(df.columns, MultiIndex)
|