fix Series.split with expand=True (#199)

* fix Series.split with expand=True

* align asterisk in split params

Author: Dr-Irv

pandas-stubs/core/indexes/base.pyi

@@ -12,6 +12,7 @@ from typing import (
 import numpy as np
 from pandas import (
     DataFrame,
+    MultiIndex,
     Series,
 )
 from pandas.core.arrays import ExtensionArray
@@ -58,7 +59,7 @@ class Index(IndexOpsMixin, PandasObject):
         tupleize_cols: bool = ...,
     ): ...
     @property
-    def str(self) -> StringMethods[Index]: ...
+    def str(self) -> StringMethods[Index, MultiIndex]: ...
     @property
     def asi8(self) -> np_ndarray_int64: ...
     def is_(self, other) -> bool: ...

pandas-stubs/core/series.pyi

@@ -854,7 +854,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     ) -> Series[S1]: ...
     def to_period(self, freq: _str | None = ..., copy: _bool = ...) -> DataFrame: ...
     @property
-    def str(self) -> StringMethods[Series]: ...
+    def str(self) -> StringMethods[Series, DataFrame]: ...
     @property
     def dt(self) -> CombinedDatetimelikeProperties: ...
     @property

pandas-stubs/core/strings.pyi

@@ -5,17 +5,25 @@ from typing import (
     Generic,
     Literal,
     Sequence,
+    TypeVar,
     overload,
 )
 
 import numpy as np
 import pandas as pd
-from pandas import Series
+from pandas import (
+    DataFrame,
+    MultiIndex,
+    Series,
+)
 from pandas.core.base import NoNewAttributesMixin
 
 from pandas._typing import T
 
-class StringMethods(NoNewAttributesMixin, Generic[T]):
+# The _TS type is what is used for the result of str.split with expand=True
+_TS = TypeVar("_TS", DataFrame, MultiIndex)
+
+class StringMethods(NoNewAttributesMixin, Generic[T, _TS]):
     def __init__(self, data: T) -> None: ...
     def __getitem__(self, key: slice | int) -> T: ...
     def __iter__(self) -> T: ...
@@ -44,11 +52,21 @@ class StringMethods(NoNewAttributesMixin, Generic[T]):
         na_rep: str | None = ...,
         join: Literal["left", "right", "outer", "inner"] = ...,
     ) -> T: ...
+    @overload
+    def split(
+        self, pat: str = ..., n: int = ..., *, expand: Literal[True], regex: bool = ...
+    ) -> _TS: ...
+    @overload
     def split(
-        self, pat: str = ..., n: int = ..., expand: bool = ..., *, regex: bool = ...
+        self, pat: str = ..., n: int = ..., *, expand: bool = ..., regex: bool = ...
     ) -> T: ...
+    @overload
+    def rsplit(
+        self, pat: str = ..., n: int = ..., *, expand: Literal[True], regex: bool = ...
+    ) -> T: ...
+    @overload
     def rsplit(
-        self, pat: str = ..., n: int = ..., expand: bool = ..., *, regex: bool = ...
+        self, pat: str = ..., n: int = ..., *, expand: bool = ..., regex: bool = ...
     ) -> T: ...
     @overload
     def partition(self, sep: str = ...) -> pd.DataFrame: ...

tests/test_indexes.py

@@ -68,3 +68,10 @@ def test_difference_none() -> None:
     # https://github.com/pandas-dev/pandas-stubs/issues/17
     ind = pd.Index([1, 2, 3])
     check(assert_type(ind.difference([1, None]), "pd.Index"), pd.Index, int)
+
+
+def test_str_split() -> None:
+    # GH 194
+    ind = pd.Index(["a-b", "c-d"])
+    check(assert_type(ind.str.split("-"), pd.Index), pd.Index)
+    check(assert_type(ind.str.split("-", expand=True), pd.MultiIndex), pd.MultiIndex)

tests/test_series.py

@@ -977,6 +977,8 @@ def test_string_accessors():
     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)
     check(assert_type(s.str.split("a"), pd.Series), pd.Series)
+    # GH 194
+    check(assert_type(s.str.split("a", expand=True), pd.DataFrame), pd.DataFrame)
     check(assert_type(s.str.startswith("a"), "pd.Series[bool]"), pd.Series, bool)
     check(assert_type(s.str.strip(), pd.Series), pd.Series)
     check(assert_type(s.str.swapcase(), pd.Series), pd.Series)