Skip to content

Commit a46eef1

Browse files
committed
Reduce duplication and dict insertion order fix
1 parent b5ead82 commit a46eef1

File tree

1 file changed

+16
-21
lines changed

1 file changed

+16
-21
lines changed

pandas/tests/reshape/test_melt.py

+16-21
Original file line numberDiff line numberDiff line change
@@ -215,33 +215,28 @@ def test_multiindex(self):
215215
@pytest.mark.parametrize("col", [
216216
pd.Series(pd.date_range('2010', periods=5, tz='US/Pacific')),
217217
pd.Series(["a", "b", "c", "a", "d"], dtype="category")])
218-
def test_pandas_dtypes_id_var(self, col):
218+
@pytest.mark.parametrize("pandas_dtype_value", [True, False])
219+
def test_pandas_dtypes_id_var(self, col, pandas_dtype_value):
219220
# GH 15785
220221
# Pandas dtype in the id
221222
df = DataFrame({'klass': range(5),
222223
'col': col,
223-
'attr1': [1, 0, 0, 0, 0],
224-
'attr2': [0, 1, 0, 0, 0]})
224+
'attr1': [1, 0, 0, 0, 0]})
225+
if pandas_dtype_value:
226+
# Pandas dtype in the value as well
227+
df['attr2'] = col
228+
expected_value = pd.concat([pd.Series([1, 0, 0, 0, 0]), col],
229+
ignore_index=True)
230+
else:
231+
df['attr2'] = [0, 1, 0, 0, 0]
232+
expected_value = [1, 0, 0, 0, 0] + [0, 1, 0, 0, 0]
225233
result = melt(df, id_vars=['klass', 'col'], var_name='attribute',
226234
value_name='value')
227-
expected = DataFrame({'klass': list(range(5)) * 2,
228-
'col': pd.concat([col] * 2, ignore_index=True),
229-
'attribute': ['attr1'] * 5 + ['attr2'] * 5,
230-
'value': [1, 0, 0, 0, 0] + [0, 1, 0, 0, 0]})
231-
tm.assert_frame_equal(result, expected)
232-
233-
# Pandas dtype in the value
234-
df = DataFrame({'klass': range(5),
235-
'col': col,
236-
'attr1': [1, 0, 0, 0, 0],
237-
'attr2': col})
238-
result = melt(df, id_vars=['klass', 'col'], var_name='attribute',
239-
value_name='value')
240-
expected = DataFrame({'klass': list(range(5)) * 2,
241-
'col': pd.concat([col] * 2, ignore_index=True),
242-
'attribute': ['attr1'] * 5 + ['attr2'] * 5,
243-
'value': pd.concat([pd.Series([1, 0, 0, 0, 0]),
244-
col], ignore_index=True)})
235+
expected = DataFrame({0: list(range(5)) * 2,
236+
1: pd.concat([col] * 2, ignore_index=True),
237+
2: ['attr1'] * 5 + ['attr2'] * 5,
238+
3: expected_value})
239+
expected.columns = ['klass', 'col', 'attribute', 'value']
245240
tm.assert_frame_equal(result, expected)
246241

247242

0 commit comments

Comments
 (0)