Add type hints and tests for str.rsplit() for expand=False

Author: pan-vlados

pandas-stubs/core/strings.pyi

@@ -79,7 +79,9 @@ class StringMethods(NoNewAttributesMixin, Generic[T, _TS, _TM, _TS2]):
     @overload
     def rsplit(self, pat: str = ..., *, n: int = ..., expand: Literal[True]) -> _TS: ...
     @overload
-    def rsplit(self, pat: str = ..., *, n: int = ..., expand: bool = ...) -> T: ...
+    def rsplit(
+        self, pat: str = ..., *, n: int = ..., expand: Literal[False] = ...
+    ) -> _TS2: ...
     @overload
     def partition(self, sep: str = ...) -> pd.DataFrame: ...
     @overload

tests/test_indexes.py

@@ -111,7 +111,7 @@ def test_difference_none() -> None:
 def test_str_split() -> None:
     # GH 194
     ind = pd.Index(["a-b", "c-d"])
-    check(assert_type(ind.str.split("-"), "pd.Index[list[str]]"), pd.Index)
+    check(assert_type(ind.str.split("-"), "pd.Index[list[str]]"), pd.Index, list)
     check(assert_type(ind.str.split("-", expand=True), pd.MultiIndex), pd.MultiIndex)
     check(
         assert_type(ind.str.split("-", expand=False), "pd.Index[list[str]]"),
@@ -120,6 +120,18 @@ def test_str_split() -> None:
     )
 
 
+def test_str_rsplit() -> None:
+    # GH 1074
+    ind = pd.Index(["a-b", "c-d"])
+    check(assert_type(ind.str.rsplit("-"), "pd.Index[list[str]]"), pd.Index, list)
+    check(assert_type(ind.str.rsplit("-", expand=True), pd.MultiIndex), pd.MultiIndex)
+    check(
+        assert_type(ind.str.rsplit("-", expand=False), "pd.Index[list[str]]"),
+        pd.Index,
+        list,
+    )
+
+
 def test_str_match() -> None:
     i = pd.Index(
         ["applep", "bananap", "Cherryp", "DATEp", "eGGpLANTp", "123p", "23.45p"]

tests/test_series.py

@@ -1548,8 +1548,13 @@ def test_string_accessors():
     check(assert_type(s.str.rindex("p"), pd.Series), pd.Series)
     check(assert_type(s.str.rjust(80), pd.Series), pd.Series)
     check(assert_type(s.str.rpartition("p"), pd.DataFrame), pd.DataFrame)
-    check(assert_type(s.str.rsplit("a"), pd.Series), pd.Series)
+    check(assert_type(s.str.rsplit("a"), "pd.Series[list[str]]"), pd.Series, list)
     check(assert_type(s.str.rsplit("a", expand=True), pd.DataFrame), pd.DataFrame)
+    check(
+        assert_type(s.str.rsplit("a", expand=False), "pd.Series[list[str]]"),
+        pd.Series,
+        list,
+    )
     check(assert_type(s.str.rstrip(), pd.Series), pd.Series)
     check(assert_type(s.str.slice(0, 4, 2), pd.Series), pd.Series)
     check(assert_type(s.str.slice_replace(0, 2, "XX"), pd.Series), pd.Series)