CLN/TYP: Groupby agg methods

Author: topper-123

pandas/core/groupby/generic.py

@@ -87,6 +87,28 @@
 if TYPE_CHECKING:
     from pandas.core.internals import Block
 
+_agg_template = """
+Compute %(f)s of group values.
+
+Parameters
+----------
+numeric_only : bool, default %(no)s
+    Include only float, int, boolean columns. If None, will attempt to use
+    everything, then use only numeric data.
+min_count : int, default %(mc)s
+    The required number of valid values to perform the operation. If fewer
+    than ``min_count`` non-NA values are present the result will be NA.
+
+Returns
+-------
+%(return_type)s
+    Computed %(f)s of values within each group.
+
+See Also
+--------
+%(return_type)s.groupby
+"""
+
 
 NamedAgg = namedtuple("NamedAgg", ["column", "aggfunc"])
 # TODO(typing) the return value on this callable should be any *scalar*.
@@ -813,6 +835,36 @@ def count(self) -> Series:
         )
         return self._reindex_output(result, fill_value=0)
 
+    @Substitution(f="sum", no=True, mc=0, return_type="Series")
+    @Appender(_agg_template)
+    def sum(self, numeric_only=True, min_count=0) -> Series:
+        return super().sum(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="prod", no=True, mc=0, return_type="Series")
+    @Appender(_agg_template)
+    def prod(self, numeric_only=True, min_count=0) -> Series:
+        return super().prod(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="min", no=False, mc=-1, return_type="Series")
+    @Appender(_agg_template)
+    def min(self, numeric_only=False, min_count=-1) -> Series:
+        return super().min(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="max", no=False, mc=-1, return_type="Series")
+    @Appender(_agg_template)
+    def max(self, numeric_only=False, min_count=-1) -> Series:
+        return super().max(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="first", no=False, mc=-1, return_type="Series")
+    @Appender(_agg_template)
+    def first(self, numeric_only=False, min_count=-1) -> Series:
+        return super().first(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="last", no=False, mc=-1, return_type="Series")
+    @Appender(_agg_template)
+    def last(self, numeric_only=False, min_count=-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)
@@ -1900,6 +1952,36 @@ def groupby_series(obj, col=None):
             results.index = ibase.default_index(len(results))
         return results
 
+    @Substitution(f="sum", no=True, mc=0, return_type="DataFrame")
+    @Appender(_agg_template)
+    def sum(self, numeric_only=True, min_count=0) -> DataFrame:
+        return super().sum(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="prod", no=True, mc=0, return_type="DataFrame")
+    @Appender(_agg_template)
+    def prod(self, numeric_only=True, min_count=0) -> DataFrame:
+        return super().prod(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="min", no=False, mc=-1, return_type="DataFrame")
+    @Appender(_agg_template)
+    def min(self, numeric_only=False, min_count=-1) -> DataFrame:
+        return super().min(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="max", no=False, mc=-1, return_type="DataFrame")
+    @Appender(_agg_template)
+    def max(self, numeric_only=False, min_count=-1) -> DataFrame:
+        return super().max(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="first", no=False, mc=-1, return_type="DataFrame")
+    @Appender(_agg_template)
+    def first(self, numeric_only=False, min_count=-1) -> DataFrame:
+        return super().first(numeric_only=numeric_only, min_count=min_count)
+
+    @Substitution(f="last", no=False, mc=-1, return_type="DataFrame")
+    @Appender(_agg_template)
+    def last(self, numeric_only=False, min_count=-1) -> DataFrame:
+        return super().last(numeric_only=numeric_only, min_count=min_count)
+
     boxplot = boxplot_frame_groupby
 
 

pandas/core/groupby/groupby.py

@@ -36,7 +36,6 @@ class providing the base-class of operations.
 from pandas._libs import Timestamp
 import pandas._libs.groupby as libgroupby
 from pandas._typing import FrameOrSeries, Scalar
-from pandas.compat import set_function_name
 from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import Appender, Substitution, cache_readonly, doc
@@ -945,6 +944,32 @@ def _wrap_transformed_output(self, output: Mapping[base.OutputKey, np.ndarray]):
     def _wrap_applied_output(self, keys, values, not_indexed_same: bool = False):
         raise AbstractMethodError(self)
 
+    def _agg_general(
+        self, numeric_only=True, min_count=-1, *, alias: str, npfunc: Callable
+    ):
+        self._set_group_selection()
+
+        # try a cython aggregation if we can
+        try:
+            return self._cython_agg_general(
+                how=alias, alt=npfunc, numeric_only=numeric_only, min_count=min_count,
+            )
+        except DataError:
+            pass
+        except NotImplementedError as err:
+            if "function is not implemented for this dtype" in str(
+                err
+            ) or "category dtype not supported" in str(err):
+                # raised in _get_cython_function, in some cases can
+                #  be trimmed by implementing cython funcs for more dtypes
+                pass
+            else:
+                raise
+
+        # apply a non-cython aggregation
+        result = self.aggregate(lambda x: npfunc(x, axis=self.axis))
+        return result
+
     def _cython_agg_general(
         self, how: str, alt=None, numeric_only: bool = True, min_count: int = -1
     ):
@@ -1438,105 +1463,62 @@ def size(self):
             result = self._obj_1d_constructor(result)
         return self._reindex_output(result, fill_value=0)
 
-    @classmethod
-    def _add_numeric_operations(cls):
-        """
-        Add numeric operations to the GroupBy generically.
+    def sum(self, numeric_only=True, min_count=0):
+        return self._agg_general(
+            numeric_only=numeric_only, min_count=min_count, alias="add", npfunc=np.sum
+        )
+
+    def prod(self, numeric_only=True, min_count=0):
+        return self._agg_general(
+            numeric_only=numeric_only, min_count=min_count, alias="prod", npfunc=np.prod
+        )
+
+    def min(self, numeric_only=False, min_count=-1):
+        return self._agg_general(
+            numeric_only=numeric_only, min_count=min_count, alias="min", npfunc=np.min
+        )
+
+    def max(self, numeric_only=False, min_count=-1):
+        return self._agg_general(
+            numeric_only=numeric_only, min_count=min_count, alias="max", npfunc=np.max
+        )
+
+    @staticmethod
+    def _get_loc(x, axis: int = 0, *, loc: int):
+        """Helper function for first/last item that isn't NA.
         """
 
-        def groupby_function(
-            name: str,
-            alias: str,
-            npfunc,
-            numeric_only: bool = True,
-            min_count: int = -1,
-        ):
+        def get_loc_notna(x, loc: int):
+            x = x.to_numpy()
+            x = x[notna(x)]
+            if len(x) == 0:
+                return np.nan
+            return x[loc]
 
-            _local_template = """
-            Compute %(f)s of group values.
-
-            Parameters
-            ----------
-            numeric_only : bool, default %(no)s
-                Include only float, int, boolean columns. If None, will attempt to use
-                everything, then use only numeric data.
-            min_count : int, default %(mc)s
-                The required number of valid values to perform the operation. If fewer
-                than ``min_count`` non-NA values are present the result will be NA.
-
-            Returns
-            -------
-            Series or DataFrame
-                Computed %(f)s of values within each group.
-            """
-
-            @Substitution(name="groupby", f=name, no=numeric_only, mc=min_count)
-            @Appender(_common_see_also)
-            @Appender(_local_template)
-            def func(self, numeric_only=numeric_only, min_count=min_count):
-                self._set_group_selection()
-
-                # try a cython aggregation if we can
-                try:
-                    return self._cython_agg_general(
-                        how=alias,
-                        alt=npfunc,
-                        numeric_only=numeric_only,
-                        min_count=min_count,
-                    )
-                except DataError:
-                    pass
-                except NotImplementedError as err:
-                    if "function is not implemented for this dtype" in str(
-                        err
-                    ) or "category dtype not supported" in str(err):
-                        # raised in _get_cython_function, in some cases can
-                        #  be trimmed by implementing cython funcs for more dtypes
-                        pass
-                    else:
-                        raise
-
-                # apply a non-cython aggregation
-                result = self.aggregate(lambda x: npfunc(x, axis=self.axis))
-                return result
-
-            set_function_name(func, name, cls)
-
-            return func
-
-        def first_compat(x, axis=0):
-            def first(x):
-                x = x.to_numpy()
-
-                x = x[notna(x)]
-                if len(x) == 0:
-                    return np.nan
-                return x[0]
-
-            if isinstance(x, DataFrame):
-                return x.apply(first, axis=axis)
-            else:
-                return first(x)
-
-        def last_compat(x, axis=0):
-            def last(x):
-                x = x.to_numpy()
-                x = x[notna(x)]
-                if len(x) == 0:
-                    return np.nan
-                return x[-1]
-
-            if isinstance(x, DataFrame):
-                return x.apply(last, axis=axis)
-            else:
-                return last(x)
-
-        cls.sum = groupby_function("sum", "add", np.sum, min_count=0)
-        cls.prod = groupby_function("prod", "prod", np.prod, min_count=0)
-        cls.min = groupby_function("min", "min", np.min, numeric_only=False)
-        cls.max = groupby_function("max", "max", np.max, numeric_only=False)
-        cls.first = groupby_function("first", "first", first_compat, numeric_only=False)
-        cls.last = groupby_function("last", "last", last_compat, numeric_only=False)
+        if isinstance(x, DataFrame):
+            return x.apply(get_loc_notna, axis=axis, loc=loc)
+        else:
+            return get_loc_notna(x, loc=loc)
+
+    def first(self, numeric_only=False, min_count=-1):
+        first_compat = partial(self._get_loc, loc=0)
+
+        return self._agg_general(
+            numeric_only=numeric_only,
+            min_count=min_count,
+            alias="first",
+            npfunc=first_compat,
+        )
+
+    def last(self, numeric_only=False, min_count=-1):
+        last_compat = partial(self._get_loc, loc=-1)
+
+        return self._agg_general(
+            numeric_only=numeric_only,
+            min_count=min_count,
+            alias="last",
+            npfunc=last_compat,
+        )
 
     @Substitution(name="groupby")
     @Appender(_common_see_also)