GroupBy[Series].count() return type should be Series[int]

Author: chrisyeh96

pandas-stubs/core/groupby/groupby.pyi

@@ -176,7 +176,10 @@ class GroupBy(BaseGroupBy[NDFrameT]):
     @overload
     def all(self: GroupBy[DataFrame], skipna: bool = ...) -> DataFrame: ...
     @final
-    def count(self) -> NDFrameT: ...
+    @overload
+    def count(self: GroupBy[Series]) -> Series[int]: ...
+    @overload
+    def count(self: GroupBy[DataFrame]) -> DataFrame: ...
     @final
     def mean(
         self,

tests/test_frame.py

@@ -1025,6 +1025,7 @@ def test_types_groupby_methods() -> None:
     check(assert_type(df.groupby("col1").sum(), pd.DataFrame), pd.DataFrame)
     check(assert_type(df.groupby("col1").prod(), pd.DataFrame), pd.DataFrame)
     check(assert_type(df.groupby("col1").sample(), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.groupby("col1").count(), pd.DataFrame), pd.DataFrame)
     check(
         assert_type(df.groupby("col1").value_counts(normalize=False), "pd.Series[int]"),
         pd.Series,

tests/test_series.py

@@ -676,6 +676,11 @@ def test_types_groupby_methods() -> None:
     check(assert_type(s.groupby(level=0).idxmax(), pd.Series), pd.Series)
     check(assert_type(s.groupby(level=0).idxmin(), pd.Series), pd.Series)
 
+    s2 = pd.Series(["w", "x", "y", "z"], index=[3, 4, 3, 4], dtype=str)
+    check(
+        assert_type(s2.groupby(level=0).count(), "pd.Series[int]"), pd.Series, np.int_
+    )
+
 
 def test_groupby_result() -> None:
     # GH 142