From 75c9febc7394b0e9d976de0ed84e771c2c2cfe69 Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Fri, 14 Feb 2025 13:13:41 +0100
Subject: [PATCH] [mypyc] Optimize str.partition and str.rpartition

---
 mypyc/doc/str_operations.rst     |  2 ++
 mypyc/primitives/str_ops.py      | 18 ++++++++++++++++++
 mypyc/test-data/fixtures/ir.py   |  2 ++
 mypyc/test-data/run-strings.test | 23 ++++++++++++++++++++---
 4 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/mypyc/doc/str_operations.rst b/mypyc/doc/str_operations.rst
index 05612fc55213f..b2e632a8bbb66 100644
--- a/mypyc/doc/str_operations.rst
+++ b/mypyc/doc/str_operations.rst
@@ -33,10 +33,12 @@ Methods
 * ``s1.endswith(s2: str)``
 * ``s1.endswith(t: tuple[str, ...])``
 * ``s.join(x: Iterable)``
+* ``s.partition(sep: str)``
 * ``s.removeprefix(prefix: str)``
 * ``s.removesuffix(suffix: str)``
 * ``s.replace(old: str, new: str)``
 * ``s.replace(old: str, new: str, count: int)``
+* ``s.rpartition(sep: str)``
 * ``s.rsplit()``
 * ``s.rsplit(sep: str)``
 * ``s.rsplit(sep: str, maxsplit: int)``
diff --git a/mypyc/primitives/str_ops.py b/mypyc/primitives/str_ops.py
index 4c82fe11beec1..d95af0cb5ead7 100644
--- a/mypyc/primitives/str_ops.py
+++ b/mypyc/primitives/str_ops.py
@@ -199,6 +199,24 @@
         error_kind=ERR_NEVER,
     )
 
+# str.partition(str)
+method_op(
+    name="partition",
+    arg_types=[str_rprimitive, str_rprimitive],
+    return_type=tuple_rprimitive,
+    c_function_name="PyUnicode_Partition",
+    error_kind=ERR_MAGIC,
+)
+
+# str.rpartition(str)
+method_op(
+    name="rpartition",
+    arg_types=[str_rprimitive, str_rprimitive],
+    return_type=tuple_rprimitive,
+    c_function_name="PyUnicode_RPartition",
+    error_kind=ERR_MAGIC,
+)
+
 # str.replace(old, new)
 method_op(
     name="replace",
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
index 50b6815d46a29..1c7346791c682 100644
--- a/mypyc/test-data/fixtures/ir.py
+++ b/mypyc/test-data/fixtures/ir.py
@@ -113,6 +113,8 @@ def startswith(self, x: Union[str, Tuple[str, ...]], start: int=..., end: int=..
     def endswith(self, x: Union[str, Tuple[str, ...]], start: int=..., end: int=...) -> bool: ...
     def replace(self, old: str, new: str, maxcount: int=...) -> str: ...
     def encode(self, encoding: str=..., errors: str=...) -> bytes: ...
+    def partition(self, sep: str, /) -> Tuple[str, str, str]: ...
+    def rpartition(self, sep: str, /) -> Tuple[str, str, str]: ...
     def removeprefix(self, prefix: str, /) -> str: ...
     def removesuffix(self, suffix: str, /) -> str: ...
 
diff --git a/mypyc/test-data/run-strings.test b/mypyc/test-data/run-strings.test
index f96824a1cad01..94fcf84f085b5 100644
--- a/mypyc/test-data/run-strings.test
+++ b/mypyc/test-data/run-strings.test
@@ -71,7 +71,8 @@ assert remove_prefix_suffix('abc', 'a') == ('bc', 'abc')
 assert remove_prefix_suffix('abc', 'c') == ('abc', 'ab')
 
 [case testStringOps]
-from typing import List, Optional
+from typing import List, Optional, Tuple
+from testutil import assertRaises
 
 def do_split(s: str, sep: Optional[str] = None, max_split: Optional[int] = None) -> List[str]:
     if sep is not None:
@@ -121,11 +122,27 @@ def test_splitlines() -> None:
     assert splitlines(s_text, False) == ["This", "is", "", "some", "long", "text."]
     assert splitlines(s_text, True) == ["This\n", "is\n", "\n", "some\n", "long\n", "text.\n"]
 
+s_partition = "Some long text"
+
+def partition(s: str, sep: str) -> Tuple[str, str, str]:
+    return s.partition(sep)
+
+def rpartition(s: str, sep: str) -> Tuple[str, str, str]:
+    return s.rpartition(sep)
+
+def test_partition() -> None:
+    assert partition(s_partition, " ") == ("Some", " ", "long text")
+    assert partition(s_partition, "Hello") == ("Some long text", "", "")
+    assert rpartition(s_partition, " ") == ("Some long", " ", "text")
+    assert rpartition(s_partition, "Hello") == ("", "", "Some long text")
+    with assertRaises(ValueError, "empty separator"):
+        partition(s_partition, "")
+    with assertRaises(ValueError, "empty separator"):
+        rpartition(s_partition, "")
+
 def getitem(s: str, index: int) -> str:
     return s[index]
 
-from testutil import assertRaises
-
 s = "abc"
 
 def test_getitem() -> None: