@@ -487,53 +487,30 @@ def test_nth_multi_index_as_expected():
487487    tm .assert_frame_equal (result , expected )
488488
489489
490- def  test_groupby_head_tail ():
490+ @pytest .mark .parametrize ( 
491+     "op, n, expected_rows" , 
492+     [ 
493+         ("head" , - 1 , []), 
494+         ("head" , 0 , []), 
495+         ("head" , 1 , [0 , 2 ]), 
496+         ("head" , 7 , [0 , 1 , 2 ]), 
497+         ("tail" , - 1 , []), 
498+         ("tail" , 0 , []), 
499+         ("tail" , 1 , [1 , 2 ]), 
500+         ("tail" , 7 , [0 , 1 , 2 ]), 
501+     ], 
502+ ) 
503+ @pytest .mark .parametrize ("columns" , [None , [], ["A" ], ["B" ], ["A" , "B" ]]) 
504+ @pytest .mark .parametrize ("as_index" , [True , False ]) 
505+ def  test_groupby_head_tail (op , n , expected_rows , columns , as_index ):
491506    df  =  DataFrame ([[1 , 2 ], [1 , 4 ], [5 , 6 ]], columns = ["A" , "B" ])
492-     g_as  =  df .groupby ("A" , as_index = True )
493-     g_not_as  =  df .groupby ("A" , as_index = False )
494- 
495-     # as_index= False, much easier 
496-     tm .assert_frame_equal (df .loc [[0 , 2 ]], g_not_as .head (1 ))
497-     tm .assert_frame_equal (df .loc [[1 , 2 ]], g_not_as .tail (1 ))
498- 
499-     empty_not_as  =  DataFrame (columns = df .columns , index = Index ([], dtype = df .index .dtype ))
500-     empty_not_as ["A" ] =  empty_not_as ["A" ].astype (df .A .dtype )
501-     empty_not_as ["B" ] =  empty_not_as ["B" ].astype (df .B .dtype )
502-     tm .assert_frame_equal (empty_not_as , g_not_as .head (0 ))
503-     tm .assert_frame_equal (empty_not_as , g_not_as .tail (0 ))
504-     tm .assert_frame_equal (empty_not_as , g_not_as .head (- 1 ))
505-     tm .assert_frame_equal (empty_not_as , g_not_as .tail (- 1 ))
506- 
507-     tm .assert_frame_equal (df , g_not_as .head (7 ))  # contains all 
508-     tm .assert_frame_equal (df , g_not_as .tail (7 ))
509- 
510-     # as_index=True, (used to be different) 
511-     df_as  =  df 
512- 
513-     tm .assert_frame_equal (df_as .loc [[0 , 2 ]], g_as .head (1 ))
514-     tm .assert_frame_equal (df_as .loc [[1 , 2 ]], g_as .tail (1 ))
515- 
516-     empty_as  =  DataFrame (index = df_as .index [:0 ], columns = df .columns )
517-     empty_as ["A" ] =  empty_not_as ["A" ].astype (df .A .dtype )
518-     empty_as ["B" ] =  empty_not_as ["B" ].astype (df .B .dtype )
519-     tm .assert_frame_equal (empty_as , g_as .head (0 ))
520-     tm .assert_frame_equal (empty_as , g_as .tail (0 ))
521-     tm .assert_frame_equal (empty_as , g_as .head (- 1 ))
522-     tm .assert_frame_equal (empty_as , g_as .tail (- 1 ))
523- 
524-     tm .assert_frame_equal (df_as , g_as .head (7 ))  # contains all 
525-     tm .assert_frame_equal (df_as , g_as .tail (7 ))
526- 
527-     # test with selection 
528-     tm .assert_frame_equal (g_as [[]].head (1 ), df_as .loc [[0 , 2 ], []])
529-     tm .assert_frame_equal (g_as [["A" ]].head (1 ), df_as .loc [[0 , 2 ], ["A" ]])
530-     tm .assert_frame_equal (g_as [["B" ]].head (1 ), df_as .loc [[0 , 2 ], ["B" ]])
531-     tm .assert_frame_equal (g_as [["A" , "B" ]].head (1 ), df_as .loc [[0 , 2 ]])
532- 
533-     tm .assert_frame_equal (g_not_as [[]].head (1 ), df_as .loc [[0 , 2 ], []])
534-     tm .assert_frame_equal (g_not_as [["A" ]].head (1 ), df_as .loc [[0 , 2 ], ["A" ]])
535-     tm .assert_frame_equal (g_not_as [["B" ]].head (1 ), df_as .loc [[0 , 2 ], ["B" ]])
536-     tm .assert_frame_equal (g_not_as [["A" , "B" ]].head (1 ), df_as .loc [[0 , 2 ]])
507+     g  =  df .groupby ("A" , as_index = as_index )
508+     expected  =  df .iloc [expected_rows ]
509+     if  columns  is  not None :
510+         g  =  g [columns ]
511+         expected  =  expected [columns ]
512+     result  =  getattr (g , op )(n )
513+     tm .assert_frame_equal (result , expected )
537514
538515
539516def  test_group_selection_cache ():
0 commit comments