use S2 python/mypy#19182

Author: MarcoGorelli

pandas-stubs/_typing.pyi

@@ -824,30 +824,8 @@ MaskType: TypeAlias = Series[bool] | np_ndarray_bool | list[bool]
 
 # Scratch types for generics
 
-S1 = TypeVar(
-    "S1",
-    bound=str
-    | bytes
-    | datetime.date
-    | datetime.time
-    | bool
-    | int
-    | float
-    | complex
-    | Dtype
-    | datetime.datetime  # includes pd.Timestamp
-    | datetime.timedelta  # includes pd.Timedelta
-    | Period
-    | Interval
-    | CategoricalDtype
-    | BaseOffset
-    | list[str],
-    default=Any,
-)
-
-S2 = TypeVar(
-    "S2",
-    bound=str
+SeriesDType: TypeAlias = (
+    str
     | bytes
     | datetime.date
     | datetime.time
@@ -862,8 +840,11 @@ S2 = TypeVar(
     | Interval
     | CategoricalDtype
     | BaseOffset
-    | list[str],
+    | list[str]
 )
+S1 = TypeVar("S1", bound=SeriesDType, default=Any)
+# Like S1, but without `default=Any`.
+S2 = TypeVar("S2", bound=SeriesDType)
 
 IndexingInt: TypeAlias = (
     int | np.int_ | np.integer | np.unsignedinteger | np.signedinteger | np.int8

pandas-stubs/core/frame.pyi

@@ -136,6 +136,7 @@ from pandas._typing import (
     ScalarT,
     SequenceNotStr,
     SeriesByT,
+    SeriesDType,
     SortKind,
     StataDateFormat,
     StorageOptions,
@@ -1391,13 +1392,23 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     @overload
     def apply(
         self,
-        f: Callable[..., S2 | NAType],
+        f: Callable[..., S2],
         axis: AxisIndex = ...,
         raw: _bool = ...,
         result_type: None = ...,
         args: Any = ...,
         **kwargs: Any,
     ) -> Series[S2]: ...
+    @overload
+    def apply(
+        self,
+        f: Callable[..., SeriesDType | NAType],
+        axis: AxisIndex = ...,
+        raw: _bool = ...,
+        result_type: None = ...,
+        args: Any = ...,
+        **kwargs: Any,
+    ) -> Series: ...
     # Since non-scalar type T is not supported in Series[T],
     # we separate this overload from the above one
     @overload
@@ -1415,7 +1426,7 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     @overload
     def apply(
         self,
-        f: Callable[..., S2 | NAType],
+        f: Callable[..., S2],
         axis: Axis = ...,
         raw: _bool = ...,
         args: Any = ...,
@@ -1424,6 +1435,17 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
         **kwargs: Any,
     ) -> Series[S2]: ...
     @overload
+    def apply(
+        self,
+        f: Callable[..., SeriesDType | NAType],
+        axis: Axis = ...,
+        raw: _bool = ...,
+        args: Any = ...,
+        *,
+        result_type: Literal["expand", "reduce"],
+        **kwargs: Any,
+    ) -> Series: ...
+    @overload
     def apply(
         self,
         f: Callable[..., ListLikeExceptSeriesAndStr | Series | Mapping[Any, Any]],
@@ -1473,10 +1495,11 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
     ) -> Series: ...
 
     # apply() overloads with default result_type of None, and keyword axis=1 matters
+    # Use S2 (TypeVar without the `Any` default) to avoid MyPy issue https://github.com/python/mypy/issues/19182.
     @overload
     def apply(
         self,
-        f: Callable[..., S2 | NAType],
+        f: Callable[..., S2],
         raw: _bool = ...,
         result_type: None = ...,
         args: Any = ...,
@@ -1485,6 +1508,17 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
         **kwargs: Any,
     ) -> Series[S2]: ...
     @overload
+    def apply(
+        self,
+        f: Callable[..., SeriesDType | NAType],
+        raw: _bool = ...,
+        result_type: None = ...,
+        args: Any = ...,
+        *,
+        axis: AxisColumn,
+        **kwargs: Any,
+    ) -> Series: ...
+    @overload
     def apply(
         self,
         f: Callable[..., ListLikeExceptSeriesAndStr | Mapping[Any, Any]],