TYP: misc typing cleanups for #29116 (#35953)

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)