Skip to content

Commit 5a6ed40

Browse files
authored
CLN: Simplify groupby head/tail tests (#37702)
1 parent a2696fe commit 5a6ed40

File tree

1 file changed

+23
-46
lines changed

1 file changed

+23
-46
lines changed

pandas/tests/groupby/test_nth.py

+23-46
Original file line numberDiff line numberDiff line change
@@ -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

539516
def test_group_selection_cache():

0 commit comments

Comments
 (0)