[mypyc] Optimize str.startswith and str.endswith with tuple argument (#18678)

cdce8p · web-flow · commit 09ac3baf7c32 · 2025-02-18T13:46:32.000Z
diff --git a/mypyc/doc/str_operations.rst b/mypyc/doc/str_operations.rst
@@ -31,6 +31,7 @@ Methods
 * ``s.encode(encoding: str)``
 * ``s.encode(encoding: str, errors: str)``
 * ``s1.endswith(s2: str)``
+* ``s1.endswith(t: tuple[str, ...])``
 * ``s.join(x: Iterable)``
 * ``s.removeprefix(prefix: str)``
 * ``s.removesuffix(suffix: str)``
@@ -45,6 +46,7 @@ Methods
 * ``s.splitlines()``
 * ``s.splitlines(keepends: bool)``
 * ``s1.startswith(s2: str)``
+* ``s1.startswith(t: tuple[str, ...])``
 
 .. note::
 
diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
@@ -725,8 +725,8 @@ PyObject *CPyStr_RSplit(PyObject *str, PyObject *sep, CPyTagged max_split);
 PyObject *CPyStr_Replace(PyObject *str, PyObject *old_substr, PyObject *new_substr, CPyTagged max_replace);
 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);
+int CPyStr_Startswith(PyObject *self, PyObject *subobj);
+int 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);
diff --git a/mypyc/lib-rt/str_ops.c b/mypyc/lib-rt/str_ops.c
@@ -161,15 +161,51 @@ PyObject *CPyStr_Replace(PyObject *str, PyObject *old_substr,
     return PyUnicode_Replace(str, old_substr, new_substr, temp_max_replace);
 }
 
-bool CPyStr_Startswith(PyObject *self, PyObject *subobj) {
+int CPyStr_Startswith(PyObject *self, PyObject *subobj) {
     Py_ssize_t start = 0;
     Py_ssize_t end = PyUnicode_GET_LENGTH(self);
+    if (PyTuple_Check(subobj)) {
+        Py_ssize_t i;
+        for (i = 0; i < PyTuple_GET_SIZE(subobj); i++) {
+            PyObject *substring = PyTuple_GET_ITEM(subobj, i);
+            if (!PyUnicode_Check(substring)) {
+                PyErr_Format(PyExc_TypeError,
+                             "tuple for startswith must only contain str, "
+                             "not %.100s",
+                             Py_TYPE(substring)->tp_name);
+                return -1;
+            }
+            int result = PyUnicode_Tailmatch(self, substring, start, end, -1);
+            if (result) {
+                return 1;
+            }
+        }
+        return 0;
+    }
     return PyUnicode_Tailmatch(self, subobj, start, end, -1);
 }
 
-bool CPyStr_Endswith(PyObject *self, PyObject *subobj) {
+int CPyStr_Endswith(PyObject *self, PyObject *subobj) {
     Py_ssize_t start = 0;
     Py_ssize_t end = PyUnicode_GET_LENGTH(self);
+    if (PyTuple_Check(subobj)) {
+        Py_ssize_t i;
+        for (i = 0; i < PyTuple_GET_SIZE(subobj); i++) {
+            PyObject *substring = PyTuple_GET_ITEM(subobj, i);
+            if (!PyUnicode_Check(substring)) {
+                PyErr_Format(PyExc_TypeError,
+                             "tuple for endswith must only contain str, "
+                             "not %.100s",
+                             Py_TYPE(substring)->tp_name);
+                return -1;
+            }
+            int result = PyUnicode_Tailmatch(self, substring, start, end, 1);
+            if (result) {
+                return 1;
+            }
+        }
+        return 0;
+    }
     return PyUnicode_Tailmatch(self, subobj, start, end, 1);
 }
 
diff --git a/mypyc/primitives/str_ops.py b/mypyc/primitives/str_ops.py
@@ -15,6 +15,7 @@
     object_rprimitive,
     pointer_rprimitive,
     str_rprimitive,
+    tuple_rprimitive,
 )
 from mypyc.primitives.registry import (
     ERR_NEG_INT,
@@ -104,20 +105,42 @@
 method_op(
     name="startswith",
     arg_types=[str_rprimitive, str_rprimitive],
-    return_type=bool_rprimitive,
+    return_type=c_int_rprimitive,
     c_function_name="CPyStr_Startswith",
+    truncated_type=bool_rprimitive,
     error_kind=ERR_NEVER,
 )
 
+# str.startswith(tuple) (return -1/0/1)
+method_op(
+    name="startswith",
+    arg_types=[str_rprimitive, tuple_rprimitive],
+    return_type=c_int_rprimitive,
+    c_function_name="CPyStr_Startswith",
+    truncated_type=bool_rprimitive,
+    error_kind=ERR_NEG_INT,
+)
+
 # str.endswith(str)
 method_op(
     name="endswith",
     arg_types=[str_rprimitive, str_rprimitive],
-    return_type=bool_rprimitive,
+    return_type=c_int_rprimitive,
     c_function_name="CPyStr_Endswith",
+    truncated_type=bool_rprimitive,
     error_kind=ERR_NEVER,
 )
 
+# str.endswith(tuple) (return -1/0/1)
+method_op(
+    name="endswith",
+    arg_types=[str_rprimitive, tuple_rprimitive],
+    return_type=c_int_rprimitive,
+    c_function_name="CPyStr_Endswith",
+    truncated_type=bool_rprimitive,
+    error_kind=ERR_NEG_INT,
+)
+
 # str.removeprefix(str)
 method_op(
     name="removeprefix",
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
@@ -109,8 +109,8 @@ def strip (self, item: str) -> str: pass
     def join(self, x: Iterable[str]) -> str: pass
     def format(self, *args: Any, **kwargs: Any) -> str: ...
     def upper(self) -> str: ...
-    def startswith(self, x: str, start: int=..., end: int=...) -> bool: ...
-    def endswith(self, x: str, start: int=..., end: int=...) -> bool: ...
+    def startswith(self, x: Union[str, Tuple[str, ...]], start: int=..., end: int=...) -> bool: ...
+    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 removeprefix(self, prefix: str, /) -> str: ...
diff --git a/mypyc/test-data/irbuild-str.test b/mypyc/test-data/irbuild-str.test
@@ -137,6 +137,73 @@ L4:
 L5:
     unreachable
 
+[case testStrStartswithEndswithTuple]
+from typing import Tuple
+
+def do_startswith(s1: str, s2: Tuple[str, ...]) -> bool:
+    return s1.startswith(s2)
+
+def do_endswith(s1: str, s2: Tuple[str, ...]) -> bool:
+    return s1.endswith(s2)
+
+def do_tuple_literal_args(s1: str) -> None:
+    x = s1.startswith(("a", "b"))
+    y = s1.endswith(("a", "b"))
+[out]
+def do_startswith(s1, s2):
+    s1 :: str
+    s2 :: tuple
+    r0 :: i32
+    r1 :: bit
+    r2 :: bool
+L0:
+    r0 = CPyStr_Startswith(s1, s2)
+    r1 = r0 >= 0 :: signed
+    r2 = truncate r0: i32 to builtins.bool
+    return r2
+def do_endswith(s1, s2):
+    s1 :: str
+    s2 :: tuple
+    r0 :: i32
+    r1 :: bit
+    r2 :: bool
+L0:
+    r0 = CPyStr_Endswith(s1, s2)
+    r1 = r0 >= 0 :: signed
+    r2 = truncate r0: i32 to builtins.bool
+    return r2
+def do_tuple_literal_args(s1):
+    s1, r0, r1 :: str
+    r2 :: tuple[str, str]
+    r3 :: object
+    r4 :: i32
+    r5 :: bit
+    r6, x :: bool
+    r7, r8 :: str
+    r9 :: tuple[str, str]
+    r10 :: object
+    r11 :: i32
+    r12 :: bit
+    r13, y :: bool
+L0:
+    r0 = 'a'
+    r1 = 'b'
+    r2 = (r0, r1)
+    r3 = box(tuple[str, str], r2)
+    r4 = CPyStr_Startswith(s1, r3)
+    r5 = r4 >= 0 :: signed
+    r6 = truncate r4: i32 to builtins.bool
+    x = r6
+    r7 = 'a'
+    r8 = 'b'
+    r9 = (r7, r8)
+    r10 = box(tuple[str, str], r9)
+    r11 = CPyStr_Endswith(s1, r10)
+    r12 = r11 >= 0 :: signed
+    r13 = truncate r11: i32 to builtins.bool
+    y = r13
+    return 1
+
 [case testStrToBool]
 def is_true(x: str) -> bool:
     if x:
diff --git a/mypyc/test-data/run-strings.test b/mypyc/test-data/run-strings.test
@@ -20,12 +20,20 @@ def eq(x: str) -> int:
     return 2
 def match(x: str, y: str) -> Tuple[bool, bool]:
     return (x.startswith(y), x.endswith(y))
+def match_tuple(x: str, y: Tuple[str, ...]) -> Tuple[bool, bool]:
+    return (x.startswith(y), x.endswith(y))
+def match_tuple_literal_args(x: str, y: str, z: str) -> Tuple[bool, bool]:
+    return (x.startswith((y, z)), x.endswith((y, z)))
 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, remove_prefix_suffix
+from native import (
+    f, g, tostr, booltostr, concat, eq, match, match_tuple,
+    match_tuple_literal_args, remove_prefix_suffix
+)
 import sys
+from testutil import assertRaises
 
 assert f() == 'some string'
 assert f() is sys.intern('some string')
@@ -45,6 +53,18 @@ assert match('abc', '') == (True, True)
 assert match('abc', 'a') == (True, False)
 assert match('abc', 'c') == (False, True)
 assert match('', 'abc') == (False, False)
+assert match_tuple('abc', ('d', 'e')) == (False, False)
+assert match_tuple('abc', ('a', 'c')) == (True, True)
+assert match_tuple('abc', ('a',)) == (True, False)
+assert match_tuple('abc', ('c',)) == (False, True)
+assert match_tuple('abc', ('x', 'y', 'z')) == (False, False)
+assert match_tuple('abc', ('x', 'y', 'z', 'a', 'c')) == (True, True)
+with assertRaises(TypeError, "tuple for startswith must only contain str"):
+    assert match_tuple('abc', (None,))
+with assertRaises(TypeError, "tuple for endswith must only contain str"):
+    assert match_tuple('abc', ('a', None))
+assert match_tuple_literal_args('abc', 'z', 'a') == (True, False)
+assert match_tuple_literal_args('abc', 'z', 'c') == (False, True)
 
 assert remove_prefix_suffix('', '') == ('', '')
 assert remove_prefix_suffix('abc', 'a') == ('bc', 'abc')
