add spitlines

Author: jgirardet

tests/snippets/bytes.py

@@ -541,14 +541,25 @@
 assert b"123456789".partition(b"45") == (b"123", b"45", b"6789")
 assert b"14523456789".partition(b"45") == (b"1", b"45", b"23456789")
 a = b"14523456789".partition(bytearray(b"45"))
-assert isinstance(a[1], bytearray) 
+assert isinstance(a[1], bytearray)
 a = b"14523456789".partition(memoryview(b"45"))
 assert isinstance(a[1], memoryview)
 
 # partition
 assert b"123456789".rpartition(b"45") == (b"123", b"45", b"6789")
 assert b"14523456789".rpartition(b"45") == (b"14523", b"45", b"6789")
 a = b"14523456789".rpartition(bytearray(b"45"))
-assert isinstance(a[1], bytearray) 
+assert isinstance(a[1], bytearray)
 a = b"14523456789".rpartition(memoryview(b"45"))
 assert isinstance(a[1], memoryview)
+
+# splitlines
+assert b"ab c\n\nde fg\rkl\r\n".splitlines() == [b"ab c", b"", b"de fg", b"kl"]
+assert b"ab c\n\nde fg\rkl\r\n".splitlines(keepends=True) == [
+    b"ab c\n",
+    b"\n",
+    b"de fg\r",
+    b"kl\r\n",
+]
+assert b"".splitlines() == []
+assert b"One line\n".splitlines() == [b"One line"]

vm/src/obj/objbyteinner.rs

@@ -308,6 +308,22 @@ impl ByteInnerExpandtabsOptions {
     }
 }
 
+#[derive(FromArgs)]
+pub struct ByteInnerSplitlinesOptions {
+    #[pyarg(positional_or_keyword, optional = true)]
+    keepends: OptionalArg<PyObjectRef>,
+}
+
+impl ByteInnerSplitlinesOptions {
+    pub fn get_value(self, vm: &VirtualMachine) -> PyResult<bool> {
+        if let OptionalArg::Present(value) = self.keepends {
+            Ok(bool::try_from_object(vm, value)?)
+        } else {
+            Ok(false)
+        }
+    }
+}
+
 impl PyByteInner {
     pub fn repr(&self) -> PyResult<String> {
         let mut res = String::with_capacity(self.elements.len());
@@ -881,6 +897,54 @@ impl PyByteInner {
 
         res
     }
+
+    pub fn splitlines(
+        &self,
+        options: ByteInnerSplitlinesOptions,
+        vm: &VirtualMachine,
+    ) -> PyResult<Vec<&[u8]>> {
+        let keepends = options.get_value(vm)?;
+
+        let mut res = vec![];
+
+        if self.elements.is_empty() {
+            return Ok(vec![]);
+        }
+
+        let mut prev_index = 0;
+        let mut index = 0;
+        let keep = if keepends { 1 } else { 0 };
+        let slice = &self.elements;
+
+        while index < slice.len() {
+            match slice[index] {
+                b'\n' => {
+                    res.push(&slice[prev_index..index + keep]);
+                    index += 1;
+                    prev_index = index;
+                }
+                b'\r' => {
+                    if index + 2 <= slice.len() && slice[index + 1] == b'\n' {
+                        res.push(&slice[prev_index..index + keep + keep]);
+                        index += 2;
+                    } else {
+                        res.push(&slice[prev_index..index + keep]);
+                        index += 1;
+                    }
+                    prev_index = index;
+                }
+                _x => {
+                    if index == slice.len() - 1 {
+                        res.push(&slice[prev_index..=index]);
+                        break;
+                    }
+                    index += 1
+                }
+            }
+        }
+
+        Ok(res)
+    }
 }
 
 pub fn try_as_byte(obj: &PyObjectRef) -> Option<Vec<u8>> {

vm/src/obj/objbytes.rs

@@ -13,7 +13,8 @@ use crate::pyobject::{PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, Py
 
 use super::objbyteinner::{
     ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
-    ByteInnerPosition, ByteInnerSplitOptions, ByteInnerTranslateOptions, PyByteInner,
+    ByteInnerPosition, ByteInnerSplitOptions, ByteInnerSplitlinesOptions,
+    ByteInnerTranslateOptions, PyByteInner,
 };
 use super::objiter;
 
@@ -379,6 +380,17 @@ impl PyBytesRef {
     fn expandtabs(self, options: ByteInnerExpandtabsOptions, vm: &VirtualMachine) -> PyResult {
         Ok(vm.ctx.new_bytes(self.inner.expandtabs(options)))
     }
+
+    #[pymethod(name = "splitlines")]
+    fn splitlines(self, options: ByteInnerSplitlinesOptions, vm: &VirtualMachine) -> PyResult {
+        let as_bytes = self
+            .inner
+            .splitlines(options, vm)?
+            .iter()
+            .map(|x| vm.ctx.new_bytes(x.to_vec()))
+            .collect::<Vec<PyObjectRef>>();
+        Ok(vm.ctx.new_list(as_bytes))
+    }
 }
 
 #[derive(Debug)]