TYP: misc typing cleanups for pandas-dev#29116

Author: simonjayhawkins

pandas/core/aggregation.py

@@ -10,21 +10,22 @@
     Callable,
     DefaultDict,
     Dict,
+    Iterable,
     List,
     Optional,
     Sequence,
     Tuple,
     Union,
 )
 
-from pandas._typing import AggFuncType, Label
+from pandas._typing import AggFuncType, FrameOrSeries, Label
 
 from pandas.core.dtypes.common import is_dict_like, is_list_like
 
 from pandas.core.base import SpecificationError
 import pandas.core.common as com
 from pandas.core.indexes.api import Index
-from pandas.core.series import FrameOrSeriesUnion, Series
+from pandas.core.series import Series
 
 
 def reconstruct_func(
@@ -276,12 +277,13 @@ def maybe_mangle_lambdas(agg_spec: Any) -> Any:
 
 
 def relabel_result(
-    result: FrameOrSeriesUnion,
+    result: FrameOrSeries,
     func: Dict[str, List[Union[Callable, str]]],
-    columns: Tuple,
-    order: List[int],
+    columns: Iterable[Label],
+    order: Iterable[int],
 ) -> Dict[Label, Series]:
-    """Internal function to reorder result if relabelling is True for
+    """
+    Internal function to reorder result if relabelling is True for
     dataframe.agg, and return the reordered result in dict.
 
     Parameters:

pandas/core/frame.py

@@ -7415,6 +7415,12 @@ def aggregate(self, func=None, axis=0, *args, **kwargs):
         if relabeling:
             # This is to keep the order to columns occurrence unchanged, and also
             # keep the order of new columns occurrence unchanged
+
+            # For the return values of reconstruct_func, if relabeling is
+            # False, columns and order will be None.
+            assert columns is not None
+            assert order is not None
+
             result_in_dict = relabel_result(result, func, columns, order)
             result = DataFrame(result_in_dict, index=columns)