gh-548 (#551)

* chore: update function definition

* chore: added idxmax and idxmin tests

* chore: remove numeric_only from series

Author: wadhah101

pandas-stubs/core/groupby/generic.pyi

@@ -143,6 +143,8 @@ class SeriesGroupBy(GroupBy, Generic[S1]):
         legend: bool = ...,
         **kwargs,
     ) -> AxesSubplot: ...
+    def idxmax(self, axis: AxisType = ..., skipna: bool = ...) -> Series: ...
+    def idxmin(self, axis: AxisType = ..., skipna: bool = ...) -> Series: ...
 
 class _DataFrameGroupByScalar(DataFrameGroupBy):
     def __iter__(self) -> Iterator[tuple[Scalar, DataFrame]]: ...
@@ -273,10 +275,10 @@ class DataFrameGroupBy(GroupBy):
     ) -> AxesSubplot | Sequence[AxesSubplot]: ...
     def idxmax(
         self, axis: AxisType = ..., skipna: bool = ..., numeric_only: bool = ...
-    ) -> Series: ...
+    ) -> DataFrame: ...
     def idxmin(
         self, axis: AxisType = ..., skipna: bool = ..., numeric_only: bool = ...
-    ) -> Series: ...
+    ) -> DataFrame: ...
     def last(self, **kwargs) -> DataFrame: ...
     def max(self, **kwargs) -> DataFrame: ...
     def mean(self, **kwargs) -> DataFrame: ...

tests/test_frame.py

@@ -877,6 +877,8 @@ def test_types_groupby_methods() -> None:
         pd.Series,
         float,
     )
+    check(assert_type(df.groupby("col1").idxmax(), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.groupby("col1").idxmin(), pd.DataFrame), pd.DataFrame)
 
 
 def test_types_groupby_agg() -> None:

tests/test_series.py

@@ -519,6 +519,8 @@ def test_types_groupby_methods() -> None:
     check(assert_type(s.groupby(level=0).var(), "pd.Series[float]"), pd.Series, float)
     check(assert_type(s.groupby(level=0).tail(), "pd.Series[int]"), pd.Series, np.int_)
     check(assert_type(s.groupby(level=0).unique(), pd.Series), pd.Series)
+    check(assert_type(s.groupby(level=0).idxmax(), pd.Series), pd.Series)
+    check(assert_type(s.groupby(level=0).idxmin(), pd.Series), pd.Series)
 
 
 def test_types_groupby_agg() -> None: