test_reorder_levels.py 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. DataFrame,
  5. MultiIndex,
  6. )
  7. import pandas._testing as tm
  8. class TestReorderLevels:
  9. def test_reorder_levels(self, frame_or_series):
  10. index = MultiIndex(
  11. levels=[["bar"], ["one", "two", "three"], [0, 1]],
  12. codes=[[0, 0, 0, 0, 0, 0], [0, 1, 2, 0, 1, 2], [0, 1, 0, 1, 0, 1]],
  13. names=["L0", "L1", "L2"],
  14. )
  15. df = DataFrame({"A": np.arange(6), "B": np.arange(6)}, index=index)
  16. obj = tm.get_obj(df, frame_or_series)
  17. # no change, position
  18. result = obj.reorder_levels([0, 1, 2])
  19. tm.assert_equal(obj, result)
  20. # no change, labels
  21. result = obj.reorder_levels(["L0", "L1", "L2"])
  22. tm.assert_equal(obj, result)
  23. # rotate, position
  24. result = obj.reorder_levels([1, 2, 0])
  25. e_idx = MultiIndex(
  26. levels=[["one", "two", "three"], [0, 1], ["bar"]],
  27. codes=[[0, 1, 2, 0, 1, 2], [0, 1, 0, 1, 0, 1], [0, 0, 0, 0, 0, 0]],
  28. names=["L1", "L2", "L0"],
  29. )
  30. expected = DataFrame({"A": np.arange(6), "B": np.arange(6)}, index=e_idx)
  31. expected = tm.get_obj(expected, frame_or_series)
  32. tm.assert_equal(result, expected)
  33. result = obj.reorder_levels([0, 0, 0])
  34. e_idx = MultiIndex(
  35. levels=[["bar"], ["bar"], ["bar"]],
  36. codes=[[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]],
  37. names=["L0", "L0", "L0"],
  38. )
  39. expected = DataFrame({"A": np.arange(6), "B": np.arange(6)}, index=e_idx)
  40. expected = tm.get_obj(expected, frame_or_series)
  41. tm.assert_equal(result, expected)
  42. result = obj.reorder_levels(["L0", "L0", "L0"])
  43. tm.assert_equal(result, expected)
  44. def test_reorder_levels_swaplevel_equivalence(
  45. self, multiindex_year_month_day_dataframe_random_data
  46. ):
  47. ymd = multiindex_year_month_day_dataframe_random_data
  48. result = ymd.reorder_levels(["month", "day", "year"])
  49. expected = ymd.swaplevel(0, 1).swaplevel(1, 2)
  50. tm.assert_frame_equal(result, expected)
  51. result = ymd["A"].reorder_levels(["month", "day", "year"])
  52. expected = ymd["A"].swaplevel(0, 1).swaplevel(1, 2)
  53. tm.assert_series_equal(result, expected)
  54. result = ymd.T.reorder_levels(["month", "day", "year"], axis=1)
  55. expected = ymd.T.swaplevel(0, 1, axis=1).swaplevel(1, 2, axis=1)
  56. tm.assert_frame_equal(result, expected)
  57. with pytest.raises(TypeError, match="hierarchical axis"):
  58. ymd.reorder_levels([1, 2], axis=1)
  59. with pytest.raises(IndexError, match="Too many levels"):
  60. ymd.index.reorder_levels([1, 2, 3])