@@ -341,18 +341,9 @@ def test_groupby_function_rename(mframe):
341
341
assert f .__name__ == name
342
342
343
343
344
- @pytest .mark .filterwarnings ("ignore:tshift is deprecated:FutureWarning" )
345
- def test_groupby_selection_with_methods (df ):
346
- # some methods which require DatetimeIndex
347
- rng = date_range ("2014" , periods = len (df ))
348
- df .index = rng
349
-
350
- g = df .groupby (["A" ])[["C" ]]
351
- g_exp = df [["C" ]].groupby (df ["A" ])
352
- # TODO check groupby with > 1 col ?
353
-
354
- # methods which are called as .foo()
355
- methods = [
344
+ @pytest .mark .parametrize (
345
+ "method" ,
346
+ [
356
347
"count" ,
357
348
"corr" ,
358
349
"cummax" ,
@@ -370,20 +361,45 @@ def test_groupby_selection_with_methods(df):
370
361
"ffill" ,
371
362
"bfill" ,
372
363
"pct_change" ,
373
- ]
364
+ ],
365
+ )
366
+ def test_groupby_selection_with_methods (df , method ):
367
+ # some methods which require DatetimeIndex
368
+ rng = date_range ("2014" , periods = len (df ))
369
+ df .index = rng
374
370
375
- for m in methods :
376
- res = getattr ( g , m )( )
377
- exp = getattr ( g_exp , m )()
371
+ g = df . groupby ([ "A" ])[[ "C" ]]
372
+ g_exp = df [[ "C" ]]. groupby ( df [ "A" ] )
373
+ # TODO check groupby with > 1 col ?
378
374
379
- # should always be frames!
380
- tm .assert_frame_equal (res , exp )
375
+ res = getattr (g , method )()
376
+ exp = getattr (g_exp , method )()
377
+
378
+ # should always be frames!
379
+ tm .assert_frame_equal (res , exp )
380
+
381
+
382
+ @pytest .mark .filterwarnings ("ignore:tshift is deprecated:FutureWarning" )
383
+ def test_groupby_selection_tshift_raises (df ):
384
+ rng = date_range ("2014" , periods = len (df ))
385
+ df .index = rng
386
+
387
+ g = df .groupby (["A" ])[["C" ]]
381
388
382
389
# check that the index cache is cleared
383
390
with pytest .raises (ValueError , match = "Freq was not set in the index" ):
384
391
# GH#35937
385
392
g .tshift ()
386
393
394
+
395
+ def test_groupby_selection_other_methods (df ):
396
+ # some methods which require DatetimeIndex
397
+ rng = date_range ("2014" , periods = len (df ))
398
+ df .index = rng
399
+
400
+ g = df .groupby (["A" ])[["C" ]]
401
+ g_exp = df [["C" ]].groupby (df ["A" ])
402
+
387
403
# methods which aren't just .foo()
388
404
tm .assert_frame_equal (g .fillna (0 ), g_exp .fillna (0 ))
389
405
tm .assert_frame_equal (g .dtypes , g_exp .dtypes )
0 commit comments