REF: de-nest _get_cython_function (pandas-dev#29609)

Author: jbrockmendel

pandas/core/groupby/generic.py

@@ -1058,14 +1058,14 @@ def _cython_agg_blocks(
 
         return new_items, new_blocks
 
-    def _aggregate_frame(self, func, *args, **kwargs):
+    def _aggregate_frame(self, func, *args, **kwargs) -> DataFrame:
         if self.grouper.nkeys != 1:
             raise AssertionError("Number of keys must be 1")
 
         axis = self.axis
         obj = self._obj_with_exclusions
 
-        result = OrderedDict()
+        result = OrderedDict()  # type: OrderedDict
         if axis != obj._info_axis_number:
             for name, data in self:
                 fres = func(data, *args, **kwargs)

pandas/core/groupby/ops.py

@@ -201,7 +201,7 @@ def apply(self, f, data: FrameOrSeries, axis: int = 0):
                 continue
 
             # group might be modified
-            group_axes = _get_axes(group)
+            group_axes = group.axes
             res = f(group)
             if not _is_indexed_like(res, group_axes):
                 mutated = True
@@ -358,40 +358,33 @@ def _is_builtin_func(self, arg):
     def _get_cython_function(self, kind: str, how: str, values, is_numeric: bool):
 
         dtype_str = values.dtype.name
+        ftype = self._cython_functions[kind][how]
 
-        def get_func(fname):
-            # see if there is a fused-type version of function
-            # only valid for numeric
-            f = getattr(libgroupby, fname, None)
-            if f is not None and is_numeric:
-                return f
-
-            # otherwise find dtype-specific version, falling back to object
-            for dt in [dtype_str, "object"]:
-                f2 = getattr(
-                    libgroupby,
-                    "{fname}_{dtype_str}".format(fname=fname, dtype_str=dt),
-                    None,
-                )
-                if f2 is not None:
-                    return f2
-
-            if hasattr(f, "__signatures__"):
-                # inspect what fused types are implemented
-                if dtype_str == "object" and "object" not in f.__signatures__:
-                    # return None so we get a NotImplementedError below
-                    #  instead of a TypeError at runtime
-                    return None
+        # see if there is a fused-type version of function
+        # only valid for numeric
+        f = getattr(libgroupby, ftype, None)
+        if f is not None and is_numeric:
             return f
 
-        ftype = self._cython_functions[kind][how]
+        # otherwise find dtype-specific version, falling back to object
+        for dt in [dtype_str, "object"]:
+            f2 = getattr(libgroupby, f"{ftype}_{dt}", None)
+            if f2 is not None:
+                return f2
+
+        if hasattr(f, "__signatures__"):
+            # inspect what fused types are implemented
+            if dtype_str == "object" and "object" not in f.__signatures__:
+                # disallow this function so we get a NotImplementedError below
+                #  instead of a TypeError at runtime
+                f = None
 
-        func = get_func(ftype)
+        func = f
 
         if func is None:
             raise NotImplementedError(
-                "function is not implemented for this dtype: "
-                "[how->{how},dtype->{dtype_str}]".format(how=how, dtype_str=dtype_str)
+                f"function is not implemented for this dtype: "
+                f"[how->{how},dtype->{dtype_str}]"
             )
 
         return func
@@ -843,13 +836,6 @@ def agg_series(self, obj: Series, func):
         return grouper.get_result()
 
 
-def _get_axes(group):
-    if isinstance(group, Series):
-        return [group.index]
-    else:
-        return group.axes
-
-
 def _is_indexed_like(obj, axes) -> bool:
     if isinstance(obj, Series):
         if len(axes) > 1: