From efd85d33e898c168e691e88fbfd78760aec141b6 Mon Sep 17 00:00:00 2001
From: Irv Lustig <irv@princeton.com>
Date: Thu, 29 Sep 2022 18:04:33 -0400
Subject: [PATCH 1/2] make Series.apply() return Series or DataFrame based on
 callable

---
 pandas-stubs/_typing.pyi     |  8 +-------
 pandas-stubs/core/series.pyi | 17 +++++++++++++++--
 tests/test_series.py         | 30 ++++++++++++++++++++++++++----
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi
index b044efc55..709b1a55d 100644
--- a/pandas-stubs/_typing.pyi
+++ b/pandas-stubs/_typing.pyi
@@ -131,13 +131,7 @@ KeysArgType = Any
 ListLike = TypeVar("ListLike", Sequence, np.ndarray, "Series", "Index")
 ListLikeU = Union[Sequence, np.ndarray, Series, Index]
 StrLike = Union[str, np.str_]
-Scalar = Union[
-    str,
-    bytes,
-    datetime.date,
-    datetime.timedelta,
-    complex,
-]
+Scalar = Union[str, bytes, datetime.date, datetime.timedelta, complex]
 ScalarT = TypeVar("ScalarT", bound=Scalar)
 # Refine the definitions below in 3.9 to use the specialized type.
 np_ndarray_int64 = npt.NDArray[np.int64]
diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
index 30e5132ba..ae557adb0 100644
--- a/pandas-stubs/core/series.pyi
+++ b/pandas-stubs/core/series.pyi
@@ -663,9 +663,22 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
         *args,
         **kwargs,
     ) -> DataFrame: ...
+    @overload
     def apply(
-        self, func: Callable, convertDType: _bool = ..., args: tuple = ..., **kwds
-    ) -> Series | DataFrame: ...
+        self,
+        func: Callable[..., Scalar],
+        convertDType: _bool = ...,
+        args: tuple = ...,
+        **kwds,
+    ) -> Series: ...
+    @overload
+    def apply(
+        self,
+        func: Callable[..., Series],
+        convertDType: _bool = ...,
+        args: tuple = ...,
+        **kwds,
+    ) -> DataFrame: ...
     def align(
         self,
         other: DataFrame | Series,
diff --git a/tests/test_series.py b/tests/test_series.py
index b07947984..75061d886 100644
--- a/tests/test_series.py
+++ b/tests/test_series.py
@@ -374,10 +374,32 @@ def test_types_unique() -> None:
 
 
 def test_types_apply() -> None:
-    s = pd.Series([-10, 2, 2, 3, 10, 10])
-    s.apply(lambda x: x**2)
-    s.apply(np.exp)
-    s.apply(str)
+    s = pd.Series([-10, 2, 2, 3.4, 10, 10])
+
+    def square(x: float) -> float:
+        return x**2
+
+    check(assert_type(s.apply(square), pd.Series), pd.Series, float)
+    check(assert_type(s.apply(np.exp), pd.Series), pd.Series, float)
+    check(assert_type(s.apply(str), pd.Series), pd.Series, str)
+
+    def makeseries(x: float) -> pd.Series:
+        return pd.Series([x, 2 * x])
+
+    check(assert_type(s.apply(makeseries), pd.DataFrame), pd.DataFrame)
+
+    # GH 293
+
+    def retseries(x: float) -> float:
+        return x
+
+    check(assert_type(s.apply(retseries).tolist(), list), list)
+
+    def get_depth(url: str) -> int:
+        return len(url)
+
+    ss = s.astype(str)
+    check(assert_type(ss.apply(get_depth), pd.Series), pd.Series, int)
 
 
 def test_types_element_wise_arithmetic() -> None:

From 27c8e4195b941003e4ebca5b35dbaf0535fd41e0 Mon Sep 17 00:00:00 2001
From: Irv Lustig <irv@princeton.com>
Date: Fri, 30 Sep 2022 11:09:57 -0400
Subject: [PATCH 2/2] allow Hashable as result of Callable to return Series

---
 pandas-stubs/core/series.pyi | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
index ae557adb0..6184fe833 100644
--- a/pandas-stubs/core/series.pyi
+++ b/pandas-stubs/core/series.pyi
@@ -666,7 +666,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     @overload
     def apply(
         self,
-        func: Callable[..., Scalar],
+        func: Callable[..., Hashable],
         convertDType: _bool = ...,
         args: tuple = ...,
         **kwds,