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