Skip to content

Commit 8d9be32

Browse files
GH1025 Typehint DataFrame.groupby.value_counts method (pandas-dev#1027)
* GH1025 Typehint DataFrame.groupby.value_counts method * GH1025 Remove commented code
1 parent 44be1fc commit 8d9be32

File tree

2 files changed

+40
-2
lines changed

2 files changed

+40
-2
lines changed

pandas-stubs/core/groupby/generic.pyi

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -302,7 +302,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
302302
) -> PlotAxes | Series: ... # Series[PlotAxes]
303303
@overload
304304
def value_counts(
305-
self,
305+
self: DataFrameGroupBy[ByT, Literal[True]],
306306
subset: ListLike | None = ...,
307307
normalize: Literal[False] = ...,
308308
sort: bool = ...,
@@ -311,13 +311,31 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
311311
) -> Series[int]: ...
312312
@overload
313313
def value_counts(
314-
self,
314+
self: DataFrameGroupBy[ByT, Literal[True]],
315315
subset: ListLike | None,
316316
normalize: Literal[True],
317317
sort: bool = ...,
318318
ascending: bool = ...,
319319
dropna: bool = ...,
320320
) -> Series[float]: ...
321+
@overload
322+
def value_counts(
323+
self: DataFrameGroupBy[ByT, Literal[False]],
324+
subset: ListLike | None = ...,
325+
normalize: Literal[False] = ...,
326+
sort: bool = ...,
327+
ascending: bool = ...,
328+
dropna: bool = ...,
329+
) -> DataFrame: ...
330+
@overload
331+
def value_counts(
332+
self: DataFrameGroupBy[ByT, Literal[False]],
333+
subset: ListLike | None,
334+
normalize: Literal[True],
335+
sort: bool = ...,
336+
ascending: bool = ...,
337+
dropna: bool = ...,
338+
) -> DataFrame: ...
321339
def take(
322340
self, indices: TakeIndexer, axis: Axis | None | NoDefault = ..., **kwargs
323341
) -> DataFrame: ...

tests/test_frame.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1026,6 +1026,7 @@ def test_types_pivot_table() -> None:
10261026

10271027

10281028
def test_types_groupby_as_index() -> None:
1029+
"""Test type of groupby.size method depending on `as_index`."""
10291030
df = pd.DataFrame({"a": [1, 2, 3]})
10301031
check(
10311032
assert_type(
@@ -1043,6 +1044,25 @@ def test_types_groupby_as_index() -> None:
10431044
)
10441045

10451046

1047+
def test_types_groupby_as_index_value_counts() -> None:
1048+
"""Test type of groupby.value_counts method depending on `as_index`."""
1049+
df = pd.DataFrame({"a": [1, 2, 3]})
1050+
check(
1051+
assert_type(
1052+
df.groupby("a", as_index=False).value_counts(),
1053+
pd.DataFrame,
1054+
),
1055+
pd.DataFrame,
1056+
)
1057+
check(
1058+
assert_type(
1059+
df.groupby("a", as_index=True).value_counts(),
1060+
"pd.Series[int]",
1061+
),
1062+
pd.Series,
1063+
)
1064+
1065+
10461066
def test_types_groupby_size() -> None:
10471067
"""Test for GH886."""
10481068
data = [

0 commit comments

Comments
 (0)