@@ -215,33 +215,28 @@ def test_multiindex(self):
215
215
@pytest .mark .parametrize ("col" , [
216
216
pd .Series (pd .date_range ('2010' , periods = 5 , tz = 'US/Pacific' )),
217
217
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 ):
219
220
# GH 15785
220
221
# Pandas dtype in the id
221
222
df = DataFrame ({'klass' : range (5 ),
222
223
'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 ]
225
233
result = melt (df , id_vars = ['klass' , 'col' ], var_name = 'attribute' ,
226
234
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' ]
245
240
tm .assert_frame_equal (result , expected )
246
241
247
242
0 commit comments