[mypyc] Optimize str.partition and str.rpartition (#18702)

`PyUnicode_Partition` and `PyUnicode_RPartition` are currently still
missing from the docs but have been part of the stable API since Python
3.2.

Author: cdce8p

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)``

mypyc/primitives/str_ops.py

@@ -197,6 +197,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",

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: ...
 

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: