Skip to content

Commit

Permalink
Fix test case and move import
Browse files Browse the repository at this point in the history
  • Loading branch information
ABizzinotto committed Oct 18, 2023
1 parent 55e8a0d commit f035363
Showing 1 changed file with 6 additions and 12 deletions.
18 changes: 6 additions & 12 deletions pandas/tests/strings/test_extract.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,17 @@
import re

import numpy as np
import pyarrow as pa
import pytest

from pandas.core.dtypes.dtypes import ArrowDtype

from pandas import (
DataFrame,
Index,
MultiIndex,
Series,
_testing as tm,
)
from pandas.core.dtypes.dtypes import ArrowDtype


def test_extract_expand_kwarg_wrong_type_raises(any_string_dtype):
Expand Down Expand Up @@ -710,16 +710,10 @@ def test_extractall_same_as_extract_subject_index(any_string_dtype):
tm.assert_frame_equal(extract_one_noname, no_match_index)


@pytest.mark.parametrize(
"data, expected_dtype",
[
(Series(["abc", "ab"], dtype=ArrowDtype(pa.string())), "string[pyarrow]"),
(Series(["abc", "ab"], dtype="string"), "string[python]"),
(Series(["abc", "ab"]), "object"),
]
)
def test_extractall_preserves_dtype(data, expected_dtype):
# Ensure that when extractall is called on a series with specific dtypes set, that
# the dtype is preserved in the resulting DataFrame's column.
result = data.str.extractall("(ab)")
assert result.dtypes[0] == expected_dtype
import pyarrow as pa

result = Series(["abc", "ab"], dtype=ArrowDtype(pa.string())).str.extractall("(ab)")
assert result.dtypes[0] == "string[pyarrow]"

0 comments on commit f035363

Please sign in to comment.