Add explicit type parameters to arguments of df.apply (#953)

* Add explicit Any type parameters for arguments used in df.apply * Add any to args, kwargs and Mapping overloads
pandas-dev · Jun 29, 2024 · 63c4567 · 63c4567
1 parent 3db439d
commit 63c4567
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 39 deletions.
diff --git 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
@@ -1250,12 +1250,12 @@ class DataFrame(NDFrame, OpsMixin):
     @overload
     def apply(
         self,
-        f: Callable[..., ListLikeExceptSeriesAndStr | Series],
+        f: Callable[..., ListLikeExceptSeriesAndStr | Series[Any]],
         axis: AxisIndex = ...,
         raw: _bool = ...,
         result_type: None = ...,
-        args=...,
-        **kwargs,
+        args: Any = ...,
+        **kwargs: Any,
     ) -> DataFrame: ...
     @overload
     def apply(
@@ -1264,21 +1264,21 @@ 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
     @overload
     def apply(
         self,
-        f: Callable[..., Mapping],
+        f: Callable[..., Mapping[Any, Any]],
         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,57 +1287,59 @@ 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(
         self,
-        f: Callable[..., ListLikeExceptSeriesAndStr | Series | Mapping],
+        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,
-        f: Callable[..., ListLikeExceptSeriesAndStr | Series | Scalar | Mapping],
+        f: Callable[
+            ..., ListLikeExceptSeriesAndStr | Series[Any] | Scalar | Mapping[Any, Any]
+        ],
         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
     @overload
     def apply(
         self,
-        f: Callable[..., Series],
+        f: Callable[..., Series[Any]],
         axis: AxisIndex = ...,
         raw: _bool = ...,
-        args=...,
+        args: Any = ...,
         *,
         result_type: Literal["reduce"],
-        **kwargs,
-    ) -> Series: ...
+        **kwargs: Any,
+    ) -> Series[Any]: ...
 
     # apply() overloads with default result_type of None, and keyword axis=1 matters
     @overload
@@ -1346,45 +1348,45 @@ 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],
+        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
     @overload
     def apply(
         self,
-        f: Callable[..., Series],
+        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