From 9e26013feefba856a95ce6fee648fa2ca963873e Mon Sep 17 00:00:00 2001
From: Brock <jbrockmendel@gmail.com>
Date: Tue, 20 Apr 2021 20:38:07 -0700
Subject: [PATCH] REF/TYP: implement NDFrameApply

---
 pandas/core/apply.py | 69 ++++++++++++++++++++++++--------------------
 1 file changed, 38 insertions(+), 31 deletions(-)

diff --git a/pandas/core/apply.py b/pandas/core/apply.py
index 86cde647cc798..9a75857c2586d 100644
--- a/pandas/core/apply.py
+++ b/pandas/core/apply.py
@@ -137,14 +137,6 @@ def f(x):
         self.orig_f: AggFuncType = func
         self.f: AggFuncType = f
 
-    @property
-    def index(self) -> Index:
-        return self.obj.index
-
-    @property
-    def agg_axis(self) -> Index:
-        return self.obj._get_agg_axis(self.axis)
-
     @abc.abstractmethod
     def apply(self) -> FrameOrSeriesUnion:
         pass
@@ -163,9 +155,8 @@ def agg(self) -> FrameOrSeriesUnion | None:
         args = self.args
         kwargs = self.kwargs
 
-        result = self.maybe_apply_str()
-        if result is not None:
-            return result
+        if isinstance(arg, str):
+            return self.maybe_apply_str()
 
         if is_dict_like(arg):
             return self.agg_dict_like()
@@ -465,27 +456,20 @@ def agg_dict_like(self) -> FrameOrSeriesUnion:
 
         return result
 
-    def maybe_apply_str(self) -> FrameOrSeriesUnion | None:
+    def maybe_apply_str(self) -> FrameOrSeriesUnion:
         """
         Compute apply in case of a string.
 
         Returns
         -------
-        result: Series, DataFrame, or None
-            Result when self.f is a string, None otherwise.
+        result: Series or DataFrame
         """
+        # Caller is responsible for checking isinstance(self.f, str)
         f = self.f
-        if not isinstance(f, str):
-            return None
+        f = cast(str, f)
 
         obj = self.obj
 
-        # TODO: GH 39993 - Avoid special-casing by replacing with lambda
-        if f == "size" and isinstance(obj, ABCDataFrame):
-            # Special-cased because DataFrame.size returns a single scalar
-            value = obj.shape[self.axis]
-            return obj._constructor_sliced(value, index=self.agg_axis, name="size")
-
         # Support for `frame.transform('method')`
         # Some methods (shift, etc.) require the axis argument, others
         # don't, so inspect and insert if necessary.
@@ -587,7 +571,22 @@ def _try_aggregate_string_function(self, obj, arg: str, *args, **kwargs):
         )
 
 
-class FrameApply(Apply):
+class NDFrameApply(Apply):
+    """
+    Methods shared by FrameApply and SeriesApply but
+    not GroupByApply or ResamplerWindowApply
+    """
+
+    @property
+    def index(self) -> Index:
+        return self.obj.index
+
+    @property
+    def agg_axis(self) -> Index:
+        return self.obj._get_agg_axis(self.axis)
+
+
+class FrameApply(NDFrameApply):
     obj: DataFrame
 
     # ---------------------------------------------------------------
@@ -644,9 +643,8 @@ def apply(self) -> FrameOrSeriesUnion:
             return self.apply_empty_result()
 
         # string dispatch
-        result = self.maybe_apply_str()
-        if result is not None:
-            return result
+        if isinstance(self.f, str):
+            return self.maybe_apply_str()
 
         # ufunc
         elif isinstance(self.f, np.ufunc):
@@ -831,6 +829,16 @@ def wrap_results(self, results: ResType, res_index: Index) -> FrameOrSeriesUnion
 
         return result
 
+    def maybe_apply_str(self) -> FrameOrSeriesUnion:
+        # Caller is responsible for checking isinstance(self.f, str)
+        # TODO: GH#39993 - Avoid special-casing by replacing with lambda
+        if self.f == "size":
+            # Special-cased because DataFrame.size returns a single scalar
+            obj = self.obj
+            value = obj.shape[self.axis]
+            return obj._constructor_sliced(value, index=self.agg_axis, name="size")
+        return super().maybe_apply_str()
+
 
 class FrameRowApply(FrameApply):
     axis = 0
@@ -967,7 +975,7 @@ def infer_to_same_shape(self, results: ResType, res_index: Index) -> DataFrame:
         return result
 
 
-class SeriesApply(Apply):
+class SeriesApply(NDFrameApply):
     obj: Series
     axis = 0
 
@@ -1001,10 +1009,9 @@ def apply(self) -> FrameOrSeriesUnion:
         if result is not None:
             return result
 
-        # if we are a string, try to dispatch
-        result = self.maybe_apply_str()
-        if result is not None:
-            return result
+        if isinstance(self.f, str):
+            # if we are a string, try to dispatch
+            return self.maybe_apply_str()
 
         return self.apply_standard()