diff --git a/pandas-stubs/core/groupby/generic.pyi b/pandas-stubs/core/groupby/generic.pyi index 512599ac..8172b1b4 100644 --- a/pandas-stubs/core/groupby/generic.pyi +++ b/pandas-stubs/core/groupby/generic.pyi @@ -302,7 +302,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]): ) -> PlotAxes | Series: ... # Series[PlotAxes] @overload def value_counts( - self, + self: DataFrameGroupBy[ByT, Literal[True]], subset: ListLike | None = ..., normalize: Literal[False] = ..., sort: bool = ..., @@ -311,13 +311,31 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]): ) -> Series[int]: ... @overload def value_counts( - self, + self: DataFrameGroupBy[ByT, Literal[True]], subset: ListLike | None, normalize: Literal[True], sort: bool = ..., ascending: bool = ..., dropna: bool = ..., ) -> Series[float]: ... + @overload + def value_counts( + self: DataFrameGroupBy[ByT, Literal[False]], + subset: ListLike | None = ..., + normalize: Literal[False] = ..., + sort: bool = ..., + ascending: bool = ..., + dropna: bool = ..., + ) -> DataFrame: ... + @overload + def value_counts( + self: DataFrameGroupBy[ByT, Literal[False]], + subset: ListLike | None, + normalize: Literal[True], + sort: bool = ..., + ascending: bool = ..., + dropna: bool = ..., + ) -> DataFrame: ... def take( self, indices: TakeIndexer, axis: Axis | None | NoDefault = ..., **kwargs ) -> DataFrame: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 7a167f53..bee41a20 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -1026,6 +1026,7 @@ def test_types_pivot_table() -> None: def test_types_groupby_as_index() -> None: + """Test type of groupby.size method depending on `as_index`.""" df = pd.DataFrame({"a": [1, 2, 3]}) check( assert_type( @@ -1043,6 +1044,25 @@ def test_types_groupby_as_index() -> None: ) +def test_types_groupby_as_index_value_counts() -> None: + """Test type of groupby.value_counts method depending on `as_index`.""" + df = pd.DataFrame({"a": [1, 2, 3]}) + check( + assert_type( + df.groupby("a", as_index=False).value_counts(), + pd.DataFrame, + ), + pd.DataFrame, + ) + check( + assert_type( + df.groupby("a", as_index=True).value_counts(), + "pd.Series[int]", + ), + pd.Series, + ) + + def test_types_groupby_size() -> None: """Test for GH886.""" data = [