test_droplevel.py 1.2 KB

123456789101112131415161718192021222324252627282930313233343536
  1. import pytest
  2. from pandas import (
  3. DataFrame,
  4. Index,
  5. MultiIndex,
  6. )
  7. import pandas._testing as tm
  8. class TestDropLevel:
  9. def test_droplevel(self, frame_or_series):
  10. # GH#20342
  11. cols = MultiIndex.from_tuples(
  12. [("c", "e"), ("d", "f")], names=["level_1", "level_2"]
  13. )
  14. mi = MultiIndex.from_tuples([(1, 2), (5, 6), (9, 10)], names=["a", "b"])
  15. df = DataFrame([[3, 4], [7, 8], [11, 12]], index=mi, columns=cols)
  16. if frame_or_series is not DataFrame:
  17. df = df.iloc[:, 0]
  18. # test that dropping of a level in index works
  19. expected = df.reset_index("a", drop=True)
  20. result = df.droplevel("a", axis="index")
  21. tm.assert_equal(result, expected)
  22. if frame_or_series is DataFrame:
  23. # test that dropping of a level in columns works
  24. expected = df.copy()
  25. expected.columns = Index(["c", "d"], name="level_1")
  26. result = df.droplevel("level_2", axis="columns")
  27. tm.assert_equal(result, expected)
  28. else:
  29. # test that droplevel raises ValueError on axis != 0
  30. with pytest.raises(ValueError, match="No axis named columns"):
  31. df.droplevel(1, axis="columns")