From 6c05e8cadac2d162fd154d357b01a494a1e58a76 Mon Sep 17 00:00:00 2001 From: Terji Petersen Date: Wed, 24 May 2023 07:23:29 +0100 Subject: [PATCH] TYP: func in core.apply.Apply.__init__ --- pandas/core/apply.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/pandas/core/apply.py b/pandas/core/apply.py index 020aa4e8916da..9b4a34961021e 100644 --- a/pandas/core/apply.py +++ b/pandas/core/apply.py @@ -103,11 +103,13 @@ def frame_apply( class Apply(metaclass=abc.ABCMeta): axis: AxisInt + orig_f: AggFuncType + f: AggFuncType def __init__( self, obj: AggObjType, - func, + func: AggFuncType, raw: bool, result_type: str | None, *, @@ -127,21 +129,19 @@ def __init__( self.result_type = result_type + f: AggFuncType # curry if needed - if ( - (kwargs or args) - and not isinstance(func, (np.ufunc, str)) - and not is_list_like(func) - ): + if callable(func) and (kwargs or args) and not isinstance(func, np.ufunc): def f(x): + assert callable(func) # needed for mypy return func(x, *args, **kwargs) else: f = func - self.orig_f: AggFuncType = func - self.f: AggFuncType = f + self.orig_f = func + self.f = f @abc.abstractmethod def apply(self) -> DataFrame | Series: @@ -1201,7 +1201,7 @@ def transform(self): def reconstruct_func( func: AggFuncType | None, **kwargs -) -> tuple[bool, AggFuncType | None, list[str] | None, npt.NDArray[np.intp] | None]: +) -> tuple[bool, AggFuncType, list[str] | None, npt.NDArray[np.intp] | None]: """ This is the internal function to reconstruct func given if there is relabeling or not and also normalize the keyword to get new order of columns. @@ -1256,6 +1256,7 @@ def reconstruct_func( if relabeling: func, columns, order = normalize_keyword_aggregation(kwargs) + assert func is not None return relabeling, func, columns, order