[mypyc] Optimize str.splitlines (#18677)

https://docs.python.org/3/c-api/unicode.html#c.PyUnicode_Splitlines

Author: cdce8p

mypyc/doc/str_operations.rst

@@ -42,6 +42,8 @@ Methods
 * ``s.split()``
 * ``s.split(sep: str)``
 * ``s.split(sep: str, maxsplit: int)``
+* ``s.splitlines()``
+* ``s.splitlines(keepends: bool)``
 * ``s1.startswith(s2: str)``
 
 .. note::

mypyc/primitives/str_ops.py

@@ -163,6 +163,19 @@
         error_kind=ERR_MAGIC,
     )
 
+# str.splitlines(...)
+str_splitlines_types: list[RType] = [str_rprimitive, bool_rprimitive]
+str_splitlines_constants: list[list[tuple[int, RType]]] = [[(0, c_int_rprimitive)], []]
+for i in range(2):
+    method_op(
+        name="splitlines",
+        arg_types=str_splitlines_types[0 : i + 1],
+        return_type=list_rprimitive,
+        c_function_name="PyUnicode_Splitlines",
+        extra_int_constants=str_splitlines_constants[i],
+        error_kind=ERR_NEVER,
+    )
+
 # str.replace(old, new)
 method_op(
     name="replace",

mypyc/test-data/fixtures/ir.py

@@ -104,6 +104,7 @@ def __contains__(self, item: str) -> bool: pass
     def __iter__(self) -> Iterator[str]: ...
     def split(self, sep: Optional[str] = None, maxsplit: int = -1) -> List[str]: pass
     def rsplit(self, sep: Optional[str] = None, maxsplit: int = -1) -> List[str]: pass
+    def splitlines(self, keepends: bool = False) -> List[str]: ...
     def strip (self, item: str) -> str: pass
     def join(self, x: Iterable[str]) -> str: pass
     def format(self, *args: Any, **kwargs: Any) -> str: ...

mypyc/test-data/run-strings.test

@@ -89,6 +89,18 @@ def test_rsplit() -> None:
     assert do_rsplit(ss, " ", 1) == ["abc abcd abcde", "abcdef"]  # different to do_split
     assert do_rsplit(ss, " ", 2) == ["abc abcd", "abcde", "abcdef"]  # different to do_split
 
+def splitlines(s: str, keepends: Optional[bool] = None) -> List[str]:
+    if keepends is not None:
+        return s.splitlines(keepends)
+    return s.splitlines()
+
+s_text = "This\nis\n\nsome\nlong\ntext.\n"
+
+def test_splitlines() -> None:
+    assert splitlines(s_text) == ["This", "is", "", "some", "long", "text."]
+    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"]
+
 def getitem(s: str, index: int) -> str:
     return s[index]