Cover all argument combinations and add a comprehensive set of test

cases for apply()

Author: gandhis1

pandas-stubs/_typing.pyi

@@ -152,7 +152,9 @@ num: TypeAlias = complex
 SeriesAxisType: TypeAlias = Literal[
     "index", 0
 ]  # Restricted subset of _AxisType for series
-AxisType: TypeAlias = Literal["columns", "index", 0, 1]
+AxisTypeIndex: TypeAlias = Literal["index", 0]
+AxisTypeColumn: TypeAlias = Literal["columns", 1]
+AxisType: TypeAlias = AxisTypeIndex | AxisTypeColumn
 DtypeNp = TypeVar("DtypeNp", bound=np.dtype[np.generic])
 KeysArgType: TypeAlias = Any
 ListLike = TypeVar("ListLike", Sequence, np.ndarray, "Series", "Index")

pandas-stubs/core/frame.pyi

@@ -59,6 +59,8 @@ from pandas._typing import (
     Axes,
     Axis,
     AxisType,
+    AxisTypeColumn,
+    AxisTypeIndex,
     CalculationMethod,
     ColspaceArgType,
     CompressionOptions,
@@ -1087,11 +1089,13 @@ class DataFrame(NDFrame, OpsMixin):
         *args,
         **kwargs,
     ) -> DataFrame: ...
+
+    # First set of apply is with defaults, and also with keyword result_type
     @overload
     def apply(
         self,
         f: Callable[..., ListLikeExceptSeriesAndStr],
-        axis: AxisType = ...,
+        axis: AxisTypeIndex = ...,
         raw: _bool = ...,
         result_type: Literal[None] = ...,
         args=...,
@@ -1101,7 +1105,7 @@ class DataFrame(NDFrame, OpsMixin):
     def apply(
         self,
         f: Callable[..., Series],
-        axis: AxisType = ...,
+        axis: AxisTypeIndex = ...,
         raw: _bool = ...,
         result_type: Literal[None] = ...,
         args=...,
@@ -1111,7 +1115,7 @@ class DataFrame(NDFrame, OpsMixin):
     def apply(
         self,
         f: Callable[..., Scalar],
-        axis: AxisType = ...,
+        axis: AxisTypeIndex = ...,
         raw: _bool = ...,
         result_type: Literal[None] = ...,
         args=...,
@@ -1121,7 +1125,7 @@ class DataFrame(NDFrame, OpsMixin):
     def apply(
         self,
         f: Callable[..., ListLikeExceptSeriesAndStr],
-        axis: AxisType = ...,
+        axis: AxisTypeIndex = ...,
         raw: _bool = ...,
         args=...,
         *,
@@ -1132,7 +1136,7 @@ class DataFrame(NDFrame, OpsMixin):
     def apply(
         self,
         f: Callable[..., ListLikeExceptSeriesAndStr | Series],
-        axis: AxisType = ...,
+        axis: AxisTypeIndex = ...,
         raw: _bool = ...,
         args=...,
         *,
@@ -1143,10 +1147,113 @@ class DataFrame(NDFrame, OpsMixin):
     def apply(
         self,
         f: Callable[..., ListLikeExceptSeriesAndStr | Series],
-        axis: AxisType = ...,
+        axis: AxisTypeIndex = ...,
+        raw: _bool = ...,
+        args=...,
+        *,
+        result_type: Literal["broadcast"],
+        **kwargs,
+    ) -> DataFrame: ...
+
+    # Second set of apply is with keyword axis=1 only
+    @overload
+    def apply(
+        self,
+        f: Callable[..., ListLikeExceptSeriesAndStr],
+        raw: _bool = ...,
+        result_type: Literal[None] = ...,
+        args=...,
+        *,
+        axis: AxisTypeColumn,
+        **kwargs,
+    ) -> Series: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., Series],
+        raw: _bool = ...,
+        result_type: Literal[None] = ...,
+        args=...,
+        *,
+        axis: AxisTypeColumn,
+        **kwargs,
+    ) -> DataFrame: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., Scalar],
+        raw: _bool = ...,
+        result_type: Literal[None] = ...,
+        args=...,
+        *,
+        axis: AxisTypeColumn,
+        **kwargs,
+    ) -> Series: ...
+
+    # Third set of apply is with keyword axis=1 and keyword result_type
+    @overload
+    def apply(
+        self,
+        f: Callable[..., ListLikeExceptSeriesAndStr],
+        raw: _bool = ...,
+        args=...,
+        *,
+        axis: AxisTypeColumn = ...,
+        result_type: Literal[None] = ...,
+        **kwargs,
+    ) -> Series: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., Series],
+        raw: _bool = ...,
+        args=...,
+        *,
+        axis: AxisTypeColumn = ...,
+        result_type: Literal[None] = ...,
+        **kwargs,
+    ) -> DataFrame: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., Scalar],
+        raw: _bool = ...,
+        args=...,
+        *,
+        axis: AxisTypeColumn = ...,
+        result_type: Literal[None] = ...,
+        **kwargs,
+    ) -> Series: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., ListLikeExceptSeriesAndStr],
+        raw: _bool = ...,
+        args=...,
+        *,
+        axis: AxisTypeColumn = ...,
+        result_type: Literal["reduce"],
+        **kwargs,
+    ) -> Series: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., ListLikeExceptSeriesAndStr | Series],
+        raw: _bool = ...,
+        args=...,
+        *,
+        axis: AxisTypeColumn = ...,
+        result_type: Literal["expand"],
+        **kwargs,
+    ) -> DataFrame: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., ListLikeExceptSeriesAndStr | Series],
         raw: _bool = ...,
         args=...,
         *,
+        axis: AxisTypeColumn = ...,
         result_type: Literal["broadcast"],
         **kwargs,
     ) -> DataFrame: ...

tests/test_frame.py

@@ -569,8 +569,8 @@ def gethead(s: pd.Series, y: int) -> pd.Series:
         pd.DataFrame,
     )
     check(
-        assert_type(df.apply(returns_listlike_of_3, axis=1), pd.DataFrame),
-        pd.DataFrame,
+        assert_type(df.apply(returns_listlike_of_3, axis=1), pd.Series),
+        pd.Series,
     )
 
     # While this call works in reality, it errors in the type checker, because this should never be called
@@ -638,6 +638,30 @@ def gethead(s: pd.Series, y: int) -> pd.Series:
         pd.DataFrame,
     )
 
+    # Test various other argument combinations to ensure all overloads are supported
+    check(
+        assert_type(df.apply(returns_scalar, axis=0), pd.Series),
+        pd.Series,
+    )
+    check(
+        assert_type(df.apply(returns_scalar, axis=0, result_type=None), pd.Series),
+        pd.Series,
+    )
+    check(
+        assert_type(df.apply(returns_scalar, 0, False, None), pd.Series),
+        pd.Series,
+    )
+    check(
+        assert_type(df.apply(returns_scalar, 0, False, result_type=None), pd.Series),
+        pd.Series,
+    )
+    check(
+        assert_type(
+            df.apply(returns_scalar, 0, raw=False, result_type=None), pd.Series
+        ),
+        pd.Series,
+    )
+
 
 def test_types_applymap() -> None:
     df = pd.DataFrame(data={"col1": [2, 1], "col2": [3, 4]})