@@ -221,8 +221,8 @@ class TestGetDummies(object):
221
221
@pytest .fixture
222
222
def df (self ):
223
223
return DataFrame ({'A' : ['a' , 'b' , 'a' ],
224
- 'B' : ['b' , 'b' , 'c' ],
225
- 'C' : [1 , 2 , 3 ]})
224
+ 'B' : ['b' , 'b' , 'c' ],
225
+ 'C' : [1 , 2 , 3 ]})
226
226
227
227
@pytest .fixture (params = ['uint8' , 'int64' , np .float64 , bool , None ])
228
228
def dtype (self , request ):
@@ -295,7 +295,7 @@ def test_basic_types(self, sparse, dtype):
295
295
result = get_dummies (s_df , columns = ['a' ], sparse = sparse , dtype = dtype )
296
296
dtype_name = self .effective_dtype (dtype ).name
297
297
298
- expected_counts = { 'int64' : 1 , 'object' : 1 }
298
+ expected_counts = {'int64' : 1 , 'object' : 1 }
299
299
expected_counts [dtype_name ] = 3 + expected_counts .get (dtype_name , 0 )
300
300
301
301
expected = Series (expected_counts ).sort_values ()
@@ -404,7 +404,7 @@ def test_dataframe_dummies_prefix_str(self, df, sparse, dtype):
404
404
[3 , 1 , 0 , 0 , 1 ]],
405
405
columns = ['C' ] + bad_columns ,
406
406
dtype = self .effective_dtype (dtype ))
407
- expected ['C' ] = [1 ,2 , 3 ]
407
+ expected ['C' ] = [1 , 2 , 3 ]
408
408
assert_frame_equal (result , expected )
409
409
410
410
def test_dataframe_dummies_subset (self , df , sparse , dtype ):
@@ -414,9 +414,8 @@ def test_dataframe_dummies_subset(self, df, sparse, dtype):
414
414
'from_A_b' : [0 , 1 , 0 ],
415
415
'B' : ['b' , 'b' , 'c' ],
416
416
'C' : [1 , 2 , 3 ]})
417
- expected [['from_A_a' , 'from_A_b' ]] = expected [['from_A_a' , 'from_A_b' ]].astype (dtype )
418
- cols = ['from_A_a' , 'from_A_b' ]
419
- expected [cols ] = expected [cols ].astype (dtype )
417
+ columns = ['from_A_a' , 'from_A_b' ]
418
+ expected [columns ] = expected [columns ].astype (dtype )
420
419
assert_frame_equal (result , expected )
421
420
422
421
def test_dataframe_dummies_prefix_sep (self , df , sparse , dtype ):
@@ -460,8 +459,10 @@ def test_dataframe_dummies_prefix_dict(self, sparse, dtype):
460
459
'from_B_b' : [1 , 1 , 0 ],
461
460
'from_B_c' : [0 , 0 , 1 ],
462
461
'C' : [1 , 2 , 3 ]})
462
+
463
463
columns = ['from_A_a' , 'from_A_b' , 'from_B_b' , 'from_B_c' ]
464
- expected [columns ] = expected [columns ].astype (self .effective_dtype (dtype ))
464
+ effective_dtype = self .effective_dtype (dtype )
465
+ expected [columns ] = expected [columns ].astype (effective_dtype )
465
466
assert_frame_equal (result , expected )
466
467
467
468
def test_dataframe_dummies_with_na (self , df , sparse , dtype ):
@@ -495,8 +496,10 @@ def test_dataframe_dummies_with_categorical(self, df, sparse, dtype):
495
496
'B_c' : [0 , 0 , 1 ],
496
497
'cat_x' : [1 , 0 , 0 ],
497
498
'cat_y' : [0 , 1 , 1 ]}).sort_index (axis = 1 )
499
+
498
500
columns = ['A_a' , 'A_b' , 'B_b' , 'B_c' , 'cat_x' , 'cat_y' ]
499
- expected [columns ] = expected [columns ].astype (self .effective_dtype (dtype ))
501
+ effective_dtype = self .effective_dtype (dtype )
502
+ expected [columns ] = expected [columns ].astype (effective_dtype )
500
503
expected .sort_index (axis = 1 )
501
504
assert_frame_equal (result , expected )
502
505
@@ -508,7 +511,8 @@ def test_basic_drop_first(self, sparse, dtype):
508
511
s_series_index = Series (s_list , list ('ABC' ))
509
512
510
513
expected = DataFrame ({'b' : [0 , 1 , 0 ],
511
- 'c' : [0 , 0 , 1 ]}, dtype = self .effective_dtype (dtype ))
514
+ 'c' : [0 , 0 , 1 ]},
515
+ dtype = self .effective_dtype (dtype ))
512
516
513
517
result = get_dummies (s_list , drop_first = True ,
514
518
sparse = sparse , dtype = dtype )
0 commit comments