From 2dd518974825f56c996ab0b218f1aaa3d61879dc Mon Sep 17 00:00:00 2001 From: Jan-Eric Nitschke <47750513+JanEricNitschke@users.noreply.github.com> Date: Sat, 29 Jun 2024 15:52:25 +0200 Subject: [PATCH 1/2] Add explicit Any type parameters for arguments used in df.apply --- pandas-stubs/_typing.pyi | 4 ++-- pandas-stubs/core/frame.pyi | 18 ++++++++++-------- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index f82cc64b..baa358a8 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -490,8 +490,8 @@ Axis: TypeAlias = AxisIndex | AxisColumn DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic]) KeysArgType: TypeAlias = Any ListLike = TypeVar("ListLike", Sequence, np.ndarray, Series, Index) -ListLikeExceptSeriesAndStr = TypeVar( - "ListLikeExceptSeriesAndStr", MutableSequence, np.ndarray, tuple, Index +ListLikeExceptSeriesAndStr: TypeAlias = ( + MutableSequence[Any] | np.ndarray[Any, Any] | tuple[Any, ...] | Index[Any] ) ListLikeU: TypeAlias = Sequence | np.ndarray | Series | Index ListLikeHashable: TypeAlias = ( diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 600ebc76..36fb2ebd 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1250,7 +1250,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def apply( self, - f: Callable[..., ListLikeExceptSeriesAndStr | Series], + f: Callable[..., ListLikeExceptSeriesAndStr | Series[Any]], axis: AxisIndex = ..., raw: _bool = ..., result_type: None = ..., @@ -1272,7 +1272,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def apply( self, - f: Callable[..., Mapping], + f: Callable[..., Mapping[Any, Any]], axis: AxisIndex = ..., raw: _bool = ..., result_type: None = ..., @@ -1295,7 +1295,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def apply( self, - f: Callable[..., ListLikeExceptSeriesAndStr | Series | Mapping], + f: Callable[..., ListLikeExceptSeriesAndStr | Series[Any] | Mapping[Any, Any]], axis: Axis = ..., raw: _bool = ..., args=..., @@ -1317,7 +1317,9 @@ class DataFrame(NDFrame, OpsMixin): @overload def apply( self, - f: Callable[..., ListLikeExceptSeriesAndStr | Series | Scalar | Mapping], + f: Callable[ + ..., ListLikeExceptSeriesAndStr | Series[Any] | Scalar | Mapping[Any, Any] + ], axis: Axis = ..., raw: _bool = ..., args=..., @@ -1330,14 +1332,14 @@ class DataFrame(NDFrame, OpsMixin): @overload def apply( self, - f: Callable[..., Series], + f: Callable[..., Series[Any]], axis: AxisIndex = ..., raw: _bool = ..., args=..., *, result_type: Literal["reduce"], **kwargs, - ) -> Series: ... + ) -> Series[Any]: ... # apply() overloads with default result_type of None, and keyword axis=1 matters @overload @@ -1365,7 +1367,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def apply( self, - f: Callable[..., Series], + f: Callable[..., Series[Any]], raw: _bool = ..., result_type: None = ..., args=..., @@ -1378,7 +1380,7 @@ class DataFrame(NDFrame, OpsMixin): @overload def apply( self, - f: Callable[..., Series], + f: Callable[..., Series[Any]], raw: _bool = ..., args=..., *, From cdb7b9acb21d0a7fb52a4903831b6587700d3725 Mon Sep 17 00:00:00 2001 From: Jan-Eric Nitschke <47750513+JanEricNitschke@users.noreply.github.com> Date: Sat, 29 Jun 2024 15:55:53 +0200 Subject: [PATCH 2/2] Add any to args, kwargs and Mapping overloads --- pandas-stubs/core/frame.pyi | 58 ++++++++++++++++++------------------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 36fb2ebd..bad2471a 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1254,8 +1254,8 @@ class DataFrame(NDFrame, OpsMixin): axis: AxisIndex = ..., raw: _bool = ..., result_type: None = ..., - args=..., - **kwargs, + args: Any = ..., + **kwargs: Any, ) -> DataFrame: ... @overload def apply( @@ -1264,8 +1264,8 @@ class DataFrame(NDFrame, OpsMixin): axis: AxisIndex = ..., raw: _bool = ..., result_type: None = ..., - args=..., - **kwargs, + args: Any = ..., + **kwargs: Any, ) -> Series[S1]: ... # Since non-scalar type T is not supported in Series[T], # we separate this overload from the above one @@ -1276,9 +1276,9 @@ class DataFrame(NDFrame, OpsMixin): axis: AxisIndex = ..., raw: _bool = ..., result_type: None = ..., - args=..., - **kwargs, - ) -> Series: ... + args: Any = ..., + **kwargs: Any, + ) -> Series[Any]: ... # apply() overloads with keyword result_type, and axis does not matter @overload @@ -1287,10 +1287,10 @@ class DataFrame(NDFrame, OpsMixin): f: Callable[..., S1], axis: Axis = ..., raw: _bool = ..., - args=..., + args: Any = ..., *, result_type: Literal["expand", "reduce"], - **kwargs, + **kwargs: Any, ) -> Series[S1]: ... @overload def apply( @@ -1298,22 +1298,22 @@ class DataFrame(NDFrame, OpsMixin): f: Callable[..., ListLikeExceptSeriesAndStr | Series[Any] | Mapping[Any, Any]], axis: Axis = ..., raw: _bool = ..., - args=..., + args: Any = ..., *, result_type: Literal["expand"], - **kwargs, + **kwargs: Any, ) -> DataFrame: ... @overload def apply( self, - f: Callable[..., ListLikeExceptSeriesAndStr | Mapping], + f: Callable[..., ListLikeExceptSeriesAndStr | Mapping[Any, Any]], axis: Axis = ..., raw: _bool = ..., - args=..., + args: Any = ..., *, result_type: Literal["reduce"], - **kwargs, - ) -> Series: ... + **kwargs: Any, + ) -> Series[Any]: ... @overload def apply( self, @@ -1322,10 +1322,10 @@ class DataFrame(NDFrame, OpsMixin): ], axis: Axis = ..., raw: _bool = ..., - args=..., + args: Any = ..., *, result_type: Literal["broadcast"], - **kwargs, + **kwargs: Any, ) -> DataFrame: ... # apply() overloads with keyword result_type, and axis does matter @@ -1335,10 +1335,10 @@ class DataFrame(NDFrame, OpsMixin): f: Callable[..., Series[Any]], axis: AxisIndex = ..., raw: _bool = ..., - args=..., + args: Any = ..., *, result_type: Literal["reduce"], - **kwargs, + **kwargs: Any, ) -> Series[Any]: ... # apply() overloads with default result_type of None, and keyword axis=1 matters @@ -1348,32 +1348,32 @@ class DataFrame(NDFrame, OpsMixin): f: Callable[..., S1], raw: _bool = ..., result_type: None = ..., - args=..., + args: Any = ..., *, axis: AxisColumn, - **kwargs, + **kwargs: Any, ) -> Series[S1]: ... @overload def apply( self, - f: Callable[..., ListLikeExceptSeriesAndStr | Mapping], + f: Callable[..., ListLikeExceptSeriesAndStr | Mapping[Any, Any]], raw: _bool = ..., result_type: None = ..., - args=..., + args: Any = ..., *, axis: AxisColumn, - **kwargs, - ) -> Series: ... + **kwargs: Any, + ) -> Series[Any]: ... @overload def apply( self, f: Callable[..., Series[Any]], raw: _bool = ..., result_type: None = ..., - args=..., + args: Any = ..., *, axis: AxisColumn, - **kwargs, + **kwargs: Any, ) -> DataFrame: ... # apply() overloads with keyword axis=1 and keyword result_type @@ -1382,11 +1382,11 @@ class DataFrame(NDFrame, OpsMixin): self, f: Callable[..., Series[Any]], raw: _bool = ..., - args=..., + args: Any = ..., *, axis: AxisColumn, result_type: Literal["reduce"], - **kwargs, + **kwargs: Any, ) -> DataFrame: ... # Add spacing between apply() overloads and remaining annotations