From e59870a1b3bbc57b2571d914a99fbb26b6ac8eb7 Mon Sep 17 00:00:00 2001 From: Roy Williams Date: Fri, 30 Aug 2024 10:57:52 -0400 Subject: [PATCH 1/2] Allow str.contains to accept a regex as well as a string --- pandas-stubs/core/strings.pyi | 2 +- tests/test_series.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/strings.pyi b/pandas-stubs/core/strings.pyi index a21074dad..4f77cc5df 100644 --- a/pandas-stubs/core/strings.pyi +++ b/pandas-stubs/core/strings.pyi @@ -96,7 +96,7 @@ class StringMethods(NoNewAttributesMixin, Generic[T, _TS]): def get(self, i: int) -> T: ... def join(self, sep: str) -> T: ... def contains( - self, pat: str, case: bool = ..., flags: int = ..., na=..., regex: bool = ... + self, pat: str | re.Pattern, case: bool = ..., flags: int = ..., na=..., regex: bool = ... ) -> Series[bool]: ... def match( self, pat: str, case: bool = ..., flags: int = ..., na: Any = ... diff --git a/tests/test_series.py b/tests/test_series.py index e68028cc4..2e96e9922 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1452,6 +1452,7 @@ def test_string_accessors(): check(assert_type(s.str.cat(sep="X"), str), str) check(assert_type(s.str.center(10), pd.Series), pd.Series) check(assert_type(s.str.contains("a"), "pd.Series[bool]"), pd.Series, np.bool_) + check(assert_type(s.str.contains(re.compile(r"a")), "pd.Series[bool]"), pd.Series, np.bool_) check(assert_type(s.str.count("pp"), "pd.Series[int]"), pd.Series, np.integer) check(assert_type(s.str.decode("utf-8"), pd.Series), pd.Series) check(assert_type(s.str.encode("latin-1"), pd.Series), pd.Series) From d388bd9487629f0ba6b4f6d71174648447a373ed Mon Sep 17 00:00:00 2001 From: Roy Williams Date: Tue, 3 Sep 2024 15:26:38 -0700 Subject: [PATCH 2/2] Apply black formatting --- pandas-stubs/core/strings.pyi | 7 ++++++- tests/test_series.py | 6 +++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/pandas-stubs/core/strings.pyi b/pandas-stubs/core/strings.pyi index 4f77cc5df..2597bd258 100644 --- a/pandas-stubs/core/strings.pyi +++ b/pandas-stubs/core/strings.pyi @@ -96,7 +96,12 @@ class StringMethods(NoNewAttributesMixin, Generic[T, _TS]): def get(self, i: int) -> T: ... def join(self, sep: str) -> T: ... def contains( - self, pat: str | re.Pattern, case: bool = ..., flags: int = ..., na=..., regex: bool = ... + self, + pat: str | re.Pattern, + case: bool = ..., + flags: int = ..., + na=..., + regex: bool = ..., ) -> Series[bool]: ... def match( self, pat: str, case: bool = ..., flags: int = ..., na: Any = ... diff --git a/tests/test_series.py b/tests/test_series.py index 2e96e9922..bad97f5b1 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -1452,7 +1452,11 @@ def test_string_accessors(): check(assert_type(s.str.cat(sep="X"), str), str) check(assert_type(s.str.center(10), pd.Series), pd.Series) check(assert_type(s.str.contains("a"), "pd.Series[bool]"), pd.Series, np.bool_) - check(assert_type(s.str.contains(re.compile(r"a")), "pd.Series[bool]"), pd.Series, np.bool_) + check( + assert_type(s.str.contains(re.compile(r"a")), "pd.Series[bool]"), + pd.Series, + np.bool_, + ) check(assert_type(s.str.count("pp"), "pd.Series[int]"), pd.Series, np.integer) check(assert_type(s.str.decode("utf-8"), pd.Series), pd.Series) check(assert_type(s.str.encode("latin-1"), pd.Series), pd.Series)