fixup str.extract

Author: MarcoGorelli

pandas-stubs/core/strings.pyi

@@ -182,13 +182,9 @@ class StringMethods(
         self, pat: str, flags: int = ..., *, expand: Literal[True] = ...
     ) -> pd.DataFrame: ...
     @overload
-    def extract(
-        self, pat: str, flags: int, expand: Literal[False]
-    ) -> Series[type[object]]: ...
+    def extract(self, pat: str, flags: int, expand: Literal[False]) -> _TO: ...
     @overload
-    def extract(
-        self, pat: str, flags: int = ..., *, expand: Literal[False]
-    ) -> Series[type[object]]: ...
+    def extract(self, pat: str, flags: int = ..., *, expand: Literal[False]) -> _TO: ...
     def extractall(self, pat: str, flags: int = ...) -> pd.DataFrame: ...
     def find(self, sub: str, start: int = ..., end: int | None = ...) -> _TI: ...
     def rfind(self, sub: str, start: int = ..., end: int | None = ...) -> _TI: ...

tests/test_series.py

@@ -1571,30 +1571,6 @@ def test_categorical_codes():
     assert_type(cat.codes, "np_ndarray_int")
 
 
-def test_series_overloads_extract():
-    s = pd.Series(
-        ["appl;ep", "ban;anap", "Cherr;yp", "DATEp", "eGGp;LANTp", "12;3p", "23.45p"]
-    )
-    check(assert_type(s.str.extract(r"[ab](\d)"), pd.DataFrame), pd.DataFrame)
-    check(
-        assert_type(s.str.extract(r"[ab](\d)", expand=True), pd.DataFrame), pd.DataFrame
-    )
-    check(
-        assert_type(
-            s.str.extract(r"[ab](\d)", expand=False), "pd.Series[type[object]]"
-        ),
-        pd.Series,
-        object,
-    )
-    check(
-        assert_type(
-            s.str.extract(r"[ab](\d)", re.IGNORECASE, False), "pd.Series[type[object]]"
-        ),
-        pd.Series,
-        object,
-    )
-
-
 def test_relops() -> None:
     # GH 175
     s: str = "abc"

tests/test_string_accessors.py

@@ -349,3 +349,48 @@ def test_index_overloads_cat():
     check(
         assert_type(idx.str.cat(unknown_idx, sep=";"), "pd.Index[str]"), pd.Index, str
     )
+
+
+def test_series_overloads_extract():
+    s = pd.Series(DATA)
+    check(assert_type(s.str.extract(r"[ab](\d)"), pd.DataFrame), pd.DataFrame)
+    check(
+        assert_type(s.str.extract(r"[ab](\d)", expand=True), pd.DataFrame), pd.DataFrame
+    )
+    check(
+        assert_type(
+            s.str.extract(r"[ab](\d)", expand=False), "pd.Series[type[object]]"
+        ),
+        pd.Series,
+        object,
+    )
+    check(
+        assert_type(
+            s.str.extract(r"[ab](\d)", re.IGNORECASE, False), "pd.Series[type[object]]"
+        ),
+        pd.Series,
+        object,
+    )
+
+
+def test_index_overloads_extract():
+    idx = pd.Index(DATA)
+    check(assert_type(idx.str.extract(r"[ab](\d)"), pd.DataFrame), pd.DataFrame)
+    check(
+        assert_type(idx.str.extract(r"[ab](\d)", expand=True), pd.DataFrame),
+        pd.DataFrame,
+    )
+    check(
+        assert_type(
+            idx.str.extract(r"[ab](\d)", expand=False), "pd.Index[type[object]]"
+        ),
+        pd.Index,
+        object,
+    )
+    check(
+        assert_type(
+            idx.str.extract(r"[ab](\d)", re.IGNORECASE, False), "pd.Index[type[object]]"
+        ),
+        pd.Index,
+        object,
+    )