@@ -472,54 +472,86 @@ def test_rename():
472472 assert result .name == expected .name
473473
474474
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 )
498501
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' ,))
503502
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 ():
504524 # errors='ignore'
505525 s = Series (range (3 ), index = list ('abc' ))
506526 result = s .drop ('bc' , errors = 'ignore' )
507- assert_series_equal (result , s )
527+ tm . assert_series_equal (result , s )
508528 result = s .drop (['a' , 'd' ], errors = 'ignore' )
509529 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 )
514531
515532 # GH 8522
516533 s = Series ([2 , 3 ], index = [True , False ])
517534 assert s .index .is_object ()
518535 result = s .drop (True )
519536 expected = Series ([3 ], index = [False ])
520- assert_series_equal (result , expected )
537+ tm .assert_series_equal (result , expected )
538+
521539
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