[mypyc] Optimize str.removeprefix and str.removesuffix (#18672)

cdce8p · web-flow · commit 8fc8d26cb428 · 2025-02-13T22:18:42.000-08:00
`str.removeprefix` and `str.removesuffix` were added in Python 3.9.
diff --git a/mypyc/doc/str_operations.rst b/mypyc/doc/str_operations.rst
@@ -32,6 +32,8 @@ Methods
 * ``s.encode(encoding: str, errors: str)``
 * ``s1.endswith(s2: str)``
 * ``s.join(x: Iterable)``
+* ``s.removeprefix(prefix: str)``
+* ``s.removesuffix(suffix: str)``
 * ``s.replace(old: str, new: str)``
 * ``s.replace(old: str, new: str, count: int)``
 * ``s.split()``
diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
@@ -726,6 +726,8 @@ PyObject *CPyStr_Append(PyObject *o1, PyObject *o2);
 PyObject *CPyStr_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
 bool CPyStr_Startswith(PyObject *self, PyObject *subobj);
 bool CPyStr_Endswith(PyObject *self, PyObject *subobj);
+PyObject *CPyStr_Removeprefix(PyObject *self, PyObject *prefix);
+PyObject *CPyStr_Removesuffix(PyObject *self, PyObject *suffix);
 bool CPyStr_IsTrue(PyObject *obj);
 Py_ssize_t CPyStr_Size_size_t(PyObject *str);
 PyObject *CPy_Decode(PyObject *obj, PyObject *encoding, PyObject *errors);
diff --git a/mypyc/lib-rt/str_ops.c b/mypyc/lib-rt/str_ops.c
@@ -164,6 +164,26 @@ bool CPyStr_Endswith(PyObject *self, PyObject *subobj) {
     return PyUnicode_Tailmatch(self, subobj, start, end, 1);
 }
 
+PyObject *CPyStr_Removeprefix(PyObject *self, PyObject *prefix) {
+    Py_ssize_t end = PyUnicode_GET_LENGTH(self);
+    int match = PyUnicode_Tailmatch(self, prefix, 0, end, -1);
+    if (match) {
+        Py_ssize_t prefix_end = PyUnicode_GET_LENGTH(prefix);
+        return PyUnicode_Substring(self, prefix_end, end);
+    }
+    return Py_NewRef(self);
+}
+
+PyObject *CPyStr_Removesuffix(PyObject *self, PyObject *suffix) {
+    Py_ssize_t end = PyUnicode_GET_LENGTH(self);
+    int match = PyUnicode_Tailmatch(self, suffix, 0, end, 1);
+    if (match) {
+        Py_ssize_t suffix_end = PyUnicode_GET_LENGTH(suffix);
+        return PyUnicode_Substring(self, 0, end - suffix_end);
+    }
+    return Py_NewRef(self);
+}
+
 /* This does a dodgy attempt to append in place  */
 PyObject *CPyStr_Append(PyObject *o1, PyObject *o2) {
     PyUnicode_Append(&o1, o2);
diff --git a/mypyc/primitives/str_ops.py b/mypyc/primitives/str_ops.py
@@ -118,6 +118,24 @@
     error_kind=ERR_NEVER,
 )
 
+# str.removeprefix(str)
+method_op(
+    name="removeprefix",
+    arg_types=[str_rprimitive, str_rprimitive],
+    return_type=str_rprimitive,
+    c_function_name="CPyStr_Removeprefix",
+    error_kind=ERR_NEVER,
+)
+
+# str.removesuffix(str)
+method_op(
+    name="removesuffix",
+    arg_types=[str_rprimitive, str_rprimitive],
+    return_type=str_rprimitive,
+    c_function_name="CPyStr_Removesuffix",
+    error_kind=ERR_NEVER,
+)
+
 # str.split(...)
 str_split_types: list[RType] = [str_rprimitive, str_rprimitive, int_rprimitive]
 str_split_functions = ["PyUnicode_Split", "PyUnicode_Split", "CPyStr_Split"]
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
@@ -111,6 +111,8 @@ def startswith(self, x: str, start: int=..., end: int=...) -> bool: ...
     def endswith(self, x: 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 removeprefix(self, prefix: str, /) -> str: ...
+    def removesuffix(self, suffix: str, /) -> str: ...
 
 class float:
     def __init__(self, x: object) -> None: pass
diff --git a/mypyc/test-data/run-strings.test b/mypyc/test-data/run-strings.test
@@ -20,9 +20,11 @@ def eq(x: str) -> int:
     return 2
 def match(x: str, y: str) -> Tuple[bool, bool]:
     return (x.startswith(y), x.endswith(y))
+def remove_prefix_suffix(x: str, y: str) -> Tuple[str, str]:
+    return (x.removeprefix(y), x.removesuffix(y))
 
 [file driver.py]
-from native import f, g, tostr, booltostr, concat, eq, match
+from native import f, g, tostr, booltostr, concat, eq, match, remove_prefix_suffix
 import sys
 
 assert f() == 'some string'
@@ -44,6 +46,10 @@ assert match('abc', 'a') == (True, False)
 assert match('abc', 'c') == (False, True)
 assert match('', 'abc') == (False, False)
 
+assert remove_prefix_suffix('', '') == ('', '')
+assert remove_prefix_suffix('abc', 'a') == ('bc', 'abc')
+assert remove_prefix_suffix('abc', 'c') == ('abc', 'ab')
+
 [case testStringOps]
 from typing import List, Optional
 
