diff --git a/ci/deps/travis-27.yaml b/ci/deps/travis-27.yaml index 5a9e206ec2c69..03f47186f781a 100644 --- a/ci/deps/travis-27.yaml +++ b/ci/deps/travis-27.yaml @@ -41,6 +41,7 @@ dependencies: # universal - pytest - pytest-xdist + - enum34 - moto==1.3.4 - hypothesis>=3.58.0 - pip: diff --git a/pandas/tests/indexes/multi/test_enum.py b/pandas/tests/indexes/multi/test_enum.py new file mode 100644 index 0000000000000..a13dd465ca813 --- /dev/null +++ b/pandas/tests/indexes/multi/test_enum.py @@ -0,0 +1,25 @@ +import pandas as pd + +import pandas.util.testing as tm +import pandas.util._test_decorators as td + + +@td.skip_if_no('enum') +def test_enum_in_multiindex(): + # GH 21298 + # Allow use of Enums as one of the factors in a MultiIndex. + from enum import Enum + MyEnum = Enum("MyEnum", "A B") + df = pd.DataFrame(columns=pd.MultiIndex.from_product(iterables=[ + MyEnum, + [1, 2] + ])) + + exp_index_0 = pd.Index([MyEnum.A, MyEnum.B], dtype='object') + tm.assert_index_equal(df.columns.levels[0], exp_index_0) + + expected = df.copy() + df = df.append({(MyEnum.A, 1): "abc", (MyEnum.B, 2): "xyz"}, + ignore_index=True) + expected.loc[0, [(MyEnum.A, 1), (MyEnum.B, 2)]] = 'abc', 'xyz' + tm.assert_frame_equal(df, expected)