From 8bf0d93bb21202ec0388fd1fe511c98770268dcb Mon Sep 17 00:00:00 2001
From: Aviv Palivoda <palaviv@gmail.com>
Date: Fri, 24 Jan 2020 22:53:44 +0200
Subject: [PATCH] Refactor string.splitlines

---
 tests/snippets/strings.py |  4 +++-
 vm/src/obj/objstr.rs      | 28 +++++++++++++++++-----------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/tests/snippets/strings.py b/tests/snippets/strings.py
index f2202962b..e676d067f 100644
--- a/tests/snippets/strings.py
+++ b/tests/snippets/strings.py
@@ -168,7 +168,9 @@ assert not 'abcd'.startswith('', 4, 3)
 
 assert '   '.isspace()
 assert 'hello\nhallo\nHallo'.splitlines() == ['hello', 'hallo', 'Hallo']
-assert 'hello\nhallo\nHallo'.splitlines(keepends=True) == ['hello\n', 'hallo\n', 'Hallo\n']
+assert 'hello\nhallo\nHallo\n'.splitlines() == ['hello', 'hallo', 'Hallo']
+assert 'hello\nhallo\nHallo'.splitlines(keepends=True) == ['hello\n', 'hallo\n', 'Hallo']
+assert 'hello\nhallo\nHallo\n'.splitlines(keepends=True) == ['hello\n', 'hallo\n', 'Hallo\n']
 assert 'abc\t12345\txyz'.expandtabs() == 'abc     12345   xyz'
 assert '-'.join(['1', '2', '3']) == '1-2-3'
 assert 'HALLO'.isupper()
diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
index 34bdf1e14..7616e2b11 100644
--- a/vm/src/obj/objstr.rs
+++ b/vm/src/obj/objstr.rs
@@ -782,19 +782,25 @@ impl PyString {
         !self.value.is_empty() && self.value.chars().all(|c| c.is_ascii())
     }
 
-    // doesn't implement keep new line delimiter just yet
     #[pymethod]
     fn splitlines(&self, args: SplitLineArgs, vm: &VirtualMachine) -> PyObjectRef {
-        let end = if let OptionalArg::Present(true) = args.keepends {
-            "\n"
-        } else {
-            ""
-        };
-        let elements = self
-            .value
-            .split('\n')
-            .map(|e| vm.ctx.new_str(e.to_string() + end))
-            .collect();
+        let keepends = args.keepends.unwrap_or(false);
+        let mut elements = vec![];
+        let mut curr = "".to_string();
+        for ch in self.value.chars() {
+            if ch == '\n' {
+                if keepends {
+                    curr.push(ch);
+                }
+                elements.push(vm.ctx.new_str(curr.clone()));
+                curr.clear();
+            } else {
+                curr.push(ch);
+            }
+        }
+        if !curr.is_empty() {
+            elements.push(vm.ctx.new_str(curr.clone()));
+        }
         vm.ctx.new_list(elements)
     }