GH1025 Typehint DataFrame.groupby.value_counts method (#1027)

* GH1025 Typehint DataFrame.groupby.value_counts method * GH1025 Remove commented code
pandas-dev · Nov 8, 2024 · 8d9be32 · 8d9be32
1 parent 44be1fc
commit 8d9be32
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 2 deletions.
diff --git 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
@@ -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 = [