make _GroupBy generic class

Author: simonjayhawkins

pandas/core/groupby/generic.py

@@ -31,7 +31,7 @@
 
 from pandas._libs import Timestamp, lib
 from pandas._typing import FrameOrSeries
-from pandas.util._decorators import Appender, Substitution, doc
+from pandas.util._decorators import Appender, Substitution
 
 from pandas.core.dtypes.cast import (
     maybe_convert_objects,
@@ -149,7 +149,7 @@ def pinner(cls):
 
 
 @pin_whitelisted_properties(Series, base.series_apply_whitelist)
-class SeriesGroupBy(GroupBy):
+class SeriesGroupBy(GroupBy[Series]):
     _apply_whitelist = base.series_apply_whitelist
 
     def _iterate_slices(self) -> Iterable[Series]:
@@ -789,30 +789,6 @@ def count(self) -> Series:
         )
         return self._reindex_output(result, fill_value=0)
 
-    @doc(GroupBy.sum.__doc__)
-    def sum(self, numeric_only: bool = True, min_count: int = 0) -> Series:
-        return super().sum(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.prod.__doc__)
-    def prod(self, numeric_only: bool = True, min_count: int = 0) -> Series:
-        return super().prod(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.min.__doc__)
-    def min(self, numeric_only: bool = False, min_count: int = -1) -> Series:
-        return super().min(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.max.__doc__)
-    def max(self, numeric_only: bool = False, min_count: int = -1) -> Series:
-        return super().max(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.first.__doc__)
-    def first(self, numeric_only: bool = False, min_count: int = -1) -> Series:
-        return super().first(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.last.__doc__)
-    def last(self, numeric_only: bool = False, min_count: int = -1) -> Series:
-        return super().last(numeric_only=numeric_only, min_count=min_count)
-
     def _apply_to_column_groupbys(self, func):
         """ return a pass thru """
         return func(self)
@@ -837,7 +813,7 @@ def pct_change(self, periods=1, fill_method="pad", limit=None, freq=None):
 
 
 @pin_whitelisted_properties(DataFrame, base.dataframe_apply_whitelist)
-class DataFrameGroupBy(GroupBy):
+class DataFrameGroupBy(GroupBy[DataFrame]):
 
     _apply_whitelist = base.dataframe_apply_whitelist
 
@@ -1887,30 +1863,6 @@ def groupby_series(obj, col=None):
             results.index = ibase.default_index(len(results))
         return results
 
-    @doc(GroupBy.sum.__doc__)
-    def sum(self, numeric_only: bool = True, min_count: int = 0) -> DataFrame:
-        return super().sum(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.prod.__doc__)
-    def prod(self, numeric_only: bool = True, min_count: int = 0) -> DataFrame:
-        return super().prod(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.min.__doc__)
-    def min(self, numeric_only: bool = False, min_count: int = -1) -> DataFrame:
-        return super().min(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.max.__doc__)
-    def max(self, numeric_only: bool = False, min_count: int = -1) -> DataFrame:
-        return super().max(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.first.__doc__)
-    def first(self, numeric_only: bool = False, min_count: int = -1) -> DataFrame:
-        return super().first(numeric_only=numeric_only, min_count=min_count)
-
-    @doc(GroupBy.last.__doc__)
-    def last(self, numeric_only: bool = False, min_count: int = -1) -> DataFrame:
-        return super().last(numeric_only=numeric_only, min_count=min_count)
-
     boxplot = boxplot_frame_groupby
 
 

pandas/core/groupby/groupby.py

@@ -17,13 +17,15 @@ class providing the base-class of operations.
     Callable,
     Dict,
     FrozenSet,
+    Generic,
     Hashable,
     Iterable,
     List,
     Mapping,
     Optional,
     Tuple,
     Type,
+    TypeVar,
     Union,
 )
 
@@ -376,13 +378,14 @@ def _group_selection_context(groupby):
 ]
 
 
-class _GroupBy(PandasObject, SelectionMixin):
+class _GroupBy(PandasObject, SelectionMixin, Generic[FrameOrSeries]):
     _group_selection = None
     _apply_whitelist: FrozenSet[str] = frozenset()
+    obj: FrameOrSeries
 
     def __init__(
         self,
-        obj: NDFrame,
+        obj: FrameOrSeries,
         keys: Optional[_KeysArgType] = None,
         axis: int = 0,
         level=None,
@@ -1079,7 +1082,11 @@ def _apply_filter(self, indices, dropna):
         return filtered
 
 
-class GroupBy(_GroupBy):
+# We require another typevar to track operations that expand dimensions, like ohlc
+FrameOrSeries2 = TypeVar("FrameOrSeries2", bound=NDFrame)
+
+
+class GroupBy(_GroupBy[FrameOrSeries]):
     """
     Class for grouping and aggregating relational data.
 
@@ -1390,25 +1397,25 @@ def size(self):
         return self._reindex_output(result, fill_value=0)
 
     @doc(_agg_template, fname="sum", no=True, mc=0)
-    def sum(self, numeric_only: bool = True, min_count: int = 0):
+    def sum(self, numeric_only: bool = True, min_count: int = 0) -> FrameOrSeries:
         return self._agg_general(
             numeric_only=numeric_only, min_count=min_count, alias="add", npfunc=np.sum
         )
 
     @doc(_agg_template, fname="prod", no=True, mc=0)
-    def prod(self, numeric_only: bool = True, min_count: int = 0):
+    def prod(self, numeric_only: bool = True, min_count: int = 0) -> FrameOrSeries:
         return self._agg_general(
             numeric_only=numeric_only, min_count=min_count, alias="prod", npfunc=np.prod
         )
 
     @doc(_agg_template, fname="min", no=False, mc=-1)
-    def min(self, numeric_only: bool = False, min_count: int = -1):
+    def min(self, numeric_only: bool = False, min_count: int = -1) -> FrameOrSeries:
         return self._agg_general(
             numeric_only=numeric_only, min_count=min_count, alias="min", npfunc=np.min
         )
 
     @doc(_agg_template, fname="max", no=False, mc=-1)
-    def max(self, numeric_only: bool = False, min_count: int = -1):
+    def max(self, numeric_only: bool = False, min_count: int = -1) -> FrameOrSeries:
         return self._agg_general(
             numeric_only=numeric_only, min_count=min_count, alias="max", npfunc=np.max
         )
@@ -1431,7 +1438,7 @@ def get_loc_notna(x, loc: int):
             return get_loc_notna(x, loc=loc)
 
     @doc(_agg_template, fname="first", no=False, mc=-1)
-    def first(self, numeric_only: bool = False, min_count: int = -1):
+    def first(self, numeric_only: bool = False, min_count: int = -1) -> FrameOrSeries:
         first_compat = partial(self._get_loc, loc=0)
 
         return self._agg_general(
@@ -1441,8 +1448,7 @@ def first(self, numeric_only: bool = False, min_count: int = -1):
             npfunc=first_compat,
         )
 
-    @doc(_agg_template, fname="last", no=False, mc=-1)
-    def last(self, numeric_only: bool = False, min_count: int = -1):
+    def last(self, numeric_only: bool = False, min_count: int = -1) -> FrameOrSeries:
         last_compat = partial(self._get_loc, loc=-1)
 
         return self._agg_general(
@@ -2467,8 +2473,8 @@ def tail(self, n=5):
         return self._selected_obj[mask]
 
     def _reindex_output(
-        self, output: FrameOrSeries, fill_value: Scalar = np.NaN
-    ) -> FrameOrSeries:
+        self, output: FrameOrSeries2, fill_value: Scalar = np.NaN
+    ) -> FrameOrSeries2:
         """
         If we have categorical groupers, then we might want to make sure that
         we have a fully re-indexed output to the levels. This means expanding