Skip to content

Commit 6db4aee

Browse files
GH1045 PR Feedback
1 parent cb09ac2 commit 6db4aee

File tree

1 file changed

+16
-2
lines changed

1 file changed

+16
-2
lines changed

tests/test_frame.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1064,21 +1064,35 @@ def test_types_groupby_as_index() -> None:
10641064
),
10651065
pd.Series,
10661066
)
1067+
check(
1068+
assert_type(
1069+
df.groupby("a").size(),
1070+
"pd.Series[int]",
1071+
),
1072+
pd.Series,
1073+
)
10671074

10681075

10691076
def test_types_groupby_as_index_list() -> None:
10701077
"""Test type of groupby.size method depending on list of grouper GH1045."""
10711078
df = pd.DataFrame({"a": [1, 1, 2], "b": [2, 3, 2]})
10721079
check(
10731080
assert_type(
1074-
df.groupby("a", as_index=False).size(),
1081+
df.groupby(["a", "b"], as_index=False).size(),
10751082
pd.DataFrame,
10761083
),
10771084
pd.DataFrame,
10781085
)
10791086
check(
10801087
assert_type(
1081-
df.groupby("a", as_index=True).size(),
1088+
df.groupby(["a", "b"], as_index=True).size(),
1089+
"pd.Series[int]",
1090+
),
1091+
pd.Series,
1092+
)
1093+
check(
1094+
assert_type(
1095+
df.groupby(["a", "b"]).size(),
10821096
"pd.Series[int]",
10831097
),
10841098
pd.Series,

0 commit comments

Comments
 (0)