Skip to content

Commit 455a310

Browse files
author
Joshua Bradt
committed
Tested types when using MultiIndex to ensure they are lists.
Any non-list type for id_vars or value_vars now raises an error message if the DataFrame has a MultiIndex for its columns.
1 parent 7406222 commit 455a310

File tree

2 files changed

+26
-0
lines changed

2 files changed

+26
-0
lines changed

pandas/core/reshape.py

+8
Original file line numberDiff line numberDiff line change
@@ -763,6 +763,10 @@ def melt(frame, id_vars=None, value_vars=None, var_name=None,
763763
if id_vars is not None:
764764
if not is_list_like(id_vars):
765765
id_vars = [id_vars]
766+
elif (isinstance(frame.columns, MultiIndex)
767+
and not isinstance(id_vars, list)):
768+
raise TypeError('id_vars must be a list of tuples when columns'
769+
' are a MultiIndex')
766770
else:
767771
id_vars = list(id_vars)
768772
else:
@@ -771,6 +775,10 @@ def melt(frame, id_vars=None, value_vars=None, var_name=None,
771775
if value_vars is not None:
772776
if not is_list_like(value_vars):
773777
value_vars = [value_vars]
778+
elif (isinstance(frame.columns, MultiIndex)
779+
and not isinstance(value_vars, list)):
780+
raise TypeError('value_vars must be a list of tuples when'
781+
' columns are a MultiIndex')
774782
else:
775783
value_vars = list(value_vars)
776784
frame = frame.loc[:, id_vars + value_vars]

pandas/tests/test_reshape.py

+18
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,24 @@ def test_value_vars_types(self):
7070
value_vars=type_(('A', 'B')))
7171
tm.assert_frame_equal(result, expected)
7272

73+
def test_id_and_value_vars_types_with_multiindex(self):
74+
expected = DataFrame({
75+
('A', 'a'): self.df1[('A', 'a')],
76+
'CAP': ['B'] * len(self.df1),
77+
'low': ['b'] * len(self.df1),
78+
'value': self.df1[('B', 'b')],
79+
}, columns=[('A', 'a'), 'CAP', 'low', 'value'])
80+
81+
for id_vars in ([('A', 'a')], ('A', 'a')):
82+
for value_vars in ([('B', 'b')], ('B', 'b')):
83+
if isinstance(id_vars, list) and isinstance(value_vars, list):
84+
result = melt(self.df1, id_vars=id_vars,
85+
value_vars=value_vars)
86+
tm.assert_frame_equal(result, expected)
87+
else:
88+
with self.assertRaisesRegex(TypeError, r'MultiIndex'):
89+
melt(self.df1, id_vars=id_vars, value_vars=value_vars)
90+
7391
def test_custom_var_name(self):
7492
result5 = melt(self.df, var_name=self.var_name)
7593
self.assertEqual(result5.columns.tolist(), ['var', 'value'])

0 commit comments

Comments
 (0)