@@ -472,54 +472,86 @@ def test_rename():
472
472
assert result .name == expected .name
473
473
474
474
475
- def test_drop ():
476
- # unique
477
- s = Series ([1 , 2 ], index = ['one' , 'two' ])
478
- expected = Series ([1 ], index = ['one' ])
479
- result = s .drop (['two' ])
480
- assert_series_equal (result , expected )
481
- result = s .drop ('two' , axis = 'rows' )
482
- assert_series_equal (result , expected )
483
-
484
- # non-unique
485
- # GH 5248
486
- s = Series ([1 , 1 , 2 ], index = ['one' , 'two' , 'one' ])
487
- expected = Series ([1 , 2 ], index = ['one' , 'one' ])
488
- result = s .drop (['two' ], axis = 0 )
489
- assert_series_equal (result , expected )
490
- result = s .drop ('two' )
491
- assert_series_equal (result , expected )
492
-
493
- expected = Series ([1 ], index = ['two' ])
494
- result = s .drop (['one' ])
495
- assert_series_equal (result , expected )
496
- result = s .drop ('one' )
497
- assert_series_equal (result , expected )
475
+ @pytest .mark .parametrize (
476
+ 'data, index, drop_labels,'
477
+ ' axis, expected_data, expected_index' ,
478
+ [
479
+ # Unique Index
480
+ ([1 , 2 ], ['one' , 'two' ], ['two' ],
481
+ 0 , [1 ], ['one' ]),
482
+ ([1 , 2 ], ['one' , 'two' ], ['two' ],
483
+ 'rows' , [1 ], ['one' ]),
484
+ ([1 , 1 , 2 ], ['one' , 'two' , 'one' ], ['two' ],
485
+ 0 , [1 , 2 ], ['one' , 'one' ]),
486
+
487
+ # GH 5248 Non-Unique Index
488
+ ([1 , 1 , 2 ], ['one' , 'two' , 'one' ], 'two' ,
489
+ 0 , [1 , 2 ], ['one' , 'one' ]),
490
+ ([1 , 1 , 2 ], ['one' , 'two' , 'one' ], ['one' ],
491
+ 0 , [1 ], ['two' ]),
492
+ ([1 , 1 , 2 ], ['one' , 'two' , 'one' ], 'one' ,
493
+ 0 , [1 ], ['two' ])])
494
+ def test_drop_unique_and_non_unique_index (data , index , axis , drop_labels ,
495
+ expected_data , expected_index ):
496
+
497
+ s = Series (data = data , index = index )
498
+ result = s .drop (drop_labels , axis = axis )
499
+ expected = Series (data = expected_data , index = expected_index )
500
+ tm .assert_series_equal (result , expected )
498
501
499
- # single string/tuple-like
500
- s = Series (range (3 ), index = list ('abc' ))
501
- pytest .raises (KeyError , s .drop , 'bc' )
502
- pytest .raises (KeyError , s .drop , ('a' ,))
503
502
503
+ @pytest .mark .parametrize (
504
+ 'data, index, drop_labels,'
505
+ ' axis, error_type, error_desc' ,
506
+ [
507
+ # single string/tuple-like
508
+ (range (3 ), list ('abc' ), 'bc' ,
509
+ 0 , KeyError , 'not found in axis' ),
510
+
511
+ # bad axis
512
+ (range (3 ), list ('abc' ), ('a' ,),
513
+ 0 , KeyError , 'not found in axis' ),
514
+ (range (3 ), list ('abc' ), 'one' ,
515
+ 'columns' , ValueError , 'No axis named columns' )])
516
+ def test_drop_exception_raised (data , index , drop_labels ,
517
+ axis , error_type , error_desc ):
518
+
519
+ with tm .assert_raises_regex (error_type , error_desc ):
520
+ Series (data , index = index ).drop (drop_labels , axis = axis )
521
+
522
+
523
+ def test_drop_with_ignore_errors ():
504
524
# errors='ignore'
505
525
s = Series (range (3 ), index = list ('abc' ))
506
526
result = s .drop ('bc' , errors = 'ignore' )
507
- assert_series_equal (result , s )
527
+ tm . assert_series_equal (result , s )
508
528
result = s .drop (['a' , 'd' ], errors = 'ignore' )
509
529
expected = s .iloc [1 :]
510
- assert_series_equal (result , expected )
511
-
512
- # bad axis
513
- pytest .raises (ValueError , s .drop , 'one' , axis = 'columns' )
530
+ tm .assert_series_equal (result , expected )
514
531
515
532
# GH 8522
516
533
s = Series ([2 , 3 ], index = [True , False ])
517
534
assert s .index .is_object ()
518
535
result = s .drop (True )
519
536
expected = Series ([3 ], index = [False ])
520
- assert_series_equal (result , expected )
537
+ tm .assert_series_equal (result , expected )
538
+
521
539
522
- # GH 16877
523
- s = Series ([2 , 3 ], index = [0 , 1 ])
524
- with tm .assert_raises_regex (KeyError , 'not contained in axis' ):
525
- s .drop ([False , True ])
540
+ @pytest .mark .parametrize ('index' , [[1 , 2 , 3 ], [1 , 1 , 3 ]])
541
+ @pytest .mark .parametrize ('drop_labels' , [[], [1 ], [3 ]])
542
+ def test_drop_empty_list (index , drop_labels ):
543
+ # GH 21494
544
+ expected_index = [i for i in index if i not in drop_labels ]
545
+ series = pd .Series (index = index ).drop (drop_labels )
546
+ tm .assert_series_equal (series , pd .Series (index = expected_index ))
547
+
548
+
549
+ @pytest .mark .parametrize ('data, index, drop_labels' , [
550
+ (None , [1 , 2 , 3 ], [1 , 4 ]),
551
+ (None , [1 , 2 , 2 ], [1 , 4 ]),
552
+ ([2 , 3 ], [0 , 1 ], [False , True ])
553
+ ])
554
+ def test_drop_non_empty_list (data , index , drop_labels ):
555
+ # GH 21494 and GH 16877
556
+ with tm .assert_raises_regex (KeyError , 'not found in axis' ):
557
+ pd .Series (data = data , index = index ).drop (drop_labels )
0 commit comments