GH456 PR Feedback

Author: loicdiridollou

pandas-stubs/core/groupby/base.pyi

@@ -10,7 +10,7 @@ class OutputKey:
     label: Hashable
     position: int
 
-reduction_kernels: TypeAlias = Literal[
+ReductionKernelType: TypeAlias = Literal[
     "all",
     "any",
     "corrwith",
@@ -37,7 +37,7 @@ reduction_kernels: TypeAlias = Literal[
     "var",
 ]
 
-transformation_kernels: TypeAlias = Literal[
+TransformationKernelType: TypeAlias = Literal[
     "bfill",
     "cumcount",
     "cummax",
@@ -53,4 +53,4 @@ transformation_kernels: TypeAlias = Literal[
     "shift",
 ]
 
-transform_kernel_allowlist: TypeAlias = reduction_kernels | transformation_kernels
+TransformReductionListType: TypeAlias = ReductionKernelType | TransformationKernelType

pandas-stubs/core/groupby/generic.pyi

@@ -19,7 +19,7 @@ from typing import (
 from matplotlib.axes import Axes as PlotAxes
 import numpy as np
 from pandas.core.frame import DataFrame
-from pandas.core.groupby.base import transform_kernel_allowlist
+from pandas.core.groupby.base import TransformReductionListType
 from pandas.core.groupby.groupby import (
     GroupBy,
     GroupByPlot,
@@ -110,7 +110,7 @@ class SeriesGroupBy(GroupBy[Series[S1]], Generic[S1, ByT]):
     ) -> UnknownSeries: ...
     @overload
     def transform(
-        self, func: transform_kernel_allowlist, *args, **kwargs
+        self, func: TransformReductionListType, *args, **kwargs
     ) -> UnknownSeries: ...
     def filter(
         self, func: Callable | str, dropna: bool = ..., *args, **kwargs
@@ -256,7 +256,7 @@ class DataFrameGroupBy(GroupBy[DataFrame], Generic[ByT, _TT]):
     ) -> DataFrame: ...
     @overload
     def transform(
-        self, func: transform_kernel_allowlist, *args, **kwargs
+        self, func: TransformReductionListType, *args, **kwargs
     ) -> DataFrame: ...
     def filter(
         self, func: Callable, dropna: bool = ..., *args, **kwargs

tests/test_series.py

@@ -1081,7 +1081,7 @@ def test_types_groupby_agg() -> None:
 
         def sum_sr(s: pd.Series[int]) -> int:
             # type of `sum` not well inferred by mypy
-            return sum(s)
+            return s.sum()
 
         check(
             assert_type(s.groupby(level=0).agg(sum_sr), "pd.Series[int]"),
@@ -1133,7 +1133,11 @@ def func(s: pd.Series[int]) -> float:
         return s.astype(float).min()
 
     s = pd.Series([1, 2, 3, 4])
-    s.groupby([1, 1, 2, 2]).agg(lambda x: x.astype(float).min())
+    check(
+        assert_type(s.groupby([1, 1, 2, 2]).agg(func), "pd.Series[float]"),
+        pd.Series,
+        np.floating,
+    )
     check(
         assert_type(s.groupby(level=0).aggregate(func), "pd.Series[float]"),
         pd.Series,
@@ -1155,7 +1159,7 @@ def func(s: pd.Series[int]) -> float:
 
         def sum_sr(s: pd.Series[int]) -> int:
             # type of `sum` not well inferred by mypy
-            return sum(s)
+            return s.sum()
 
         check(
             assert_type(s.groupby(level=0).aggregate(sum_sr), "pd.Series[int]"),