[mypyc] Optimize str.startswith and str.endswith with tuple argument

Author: cdce8p

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)``
@@ -43,6 +44,7 @@ Methods
 * ``s.split(sep: str)``
 * ``s.split(sep: str, maxsplit: int)``
 * ``s1.startswith(s2: str)``
+* ``s1.startswith(t: tuple[str, ...])``
 
 .. note::
 

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

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);
 }
 

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,
@@ -109,6 +110,15 @@
     error_kind=ERR_NEVER,
 )
 
+# str.startswith(tuple) (return -1/0/1)
+method_op(
+    name="startswith",
+    arg_types=[str_rprimitive, tuple_rprimitive],
+    return_type=bool_rprimitive,
+    c_function_name="CPyStr_Startswith",
+    error_kind=ERR_NEG_INT,
+)
+
 # str.endswith(str)
 method_op(
     name="endswith",
@@ -118,6 +128,15 @@
     error_kind=ERR_NEVER,
 )
 
+# str.endswith(tuple) (return -1/0/1)
+method_op(
+    name="endswith",
+    arg_types=[str_rprimitive, tuple_rprimitive],
+    return_type=bool_rprimitive,
+    c_function_name="CPyStr_Endswith",
+    error_kind=ERR_NEG_INT,
+)
+
 # str.removeprefix(str)
 method_op(
     name="removeprefix",

mypyc/test-data/fixtures/ir.py

@@ -108,8 +108,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: ...

mypyc/test-data/run-strings.test

@@ -20,11 +20,13 @@ 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 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, remove_prefix_suffix
 import sys
 
 assert f() == 'some string'
@@ -45,6 +47,10 @@ 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 remove_prefix_suffix('', '') == ('', '')
 assert remove_prefix_suffix('abc', 'a') == ('bc', 'abc')