[GR-23265] Implement bytes.splitlines

PullRequest: graalpython/1066

Author: msimacek

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/bytes/ByteArrayBuiltins.java

@@ -207,6 +207,16 @@ public PNone clear(VirtualFrame frame, PByteArray byteArray,
         }
     }
 
+    // bytearray.splitlines([keepends])
+    @Builtin(name = "splitlines", minNumOfPositionalArgs = 1, parameterNames = {"self", "keepends"})
+    @GenerateNodeFactory
+    public abstract static class SplitLinesNode extends BytesBuiltins.SplitLinesNode {
+        @Override
+        protected PIBytesLike createElement(byte[] bytes) {
+            return factory().createByteArray(bytes);
+        }
+    }
+
     @Builtin(name = __SETITEM__, minNumOfPositionalArgs = 3)
     @GenerateNodeFactory
     @ImportStatic(SpecialMethodNames.class)

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/bytes/BytesBuiltins.java

@@ -90,6 +90,7 @@
 import com.oracle.graal.python.nodes.subscript.SliceLiteralNode.CastToSliceComponentNode;
 import com.oracle.graal.python.nodes.truffle.PythonArithmeticTypes;
 import com.oracle.graal.python.nodes.util.CastToByteNode;
+import com.oracle.graal.python.nodes.util.CastToJavaIntExactNode;
 import com.oracle.graal.python.runtime.exception.PythonErrorType;
 import com.oracle.graal.python.runtime.sequence.storage.ByteSequenceStorage;
 import com.oracle.graal.python.runtime.sequence.storage.IntSequenceStorage;
@@ -1645,6 +1646,61 @@ protected AbstractSplitNode createRecursiveNode() {
         }
     }
 
+    // bytes.splitlines([keepends])
+    @Builtin(name = "splitlines", minNumOfPositionalArgs = 1, parameterNames = {"self", "keepends"})
+    @GenerateNodeFactory
+    public abstract static class SplitLinesNode extends PythonBinaryBuiltinNode {
+        @Child private BytesNodes.ToBytesNode toBytesNode = BytesNodes.ToBytesNode.create();
+        @Child private AppendNode appendNode = AppendNode.create();
+
+        @Specialization
+        PList doSplitlinesDefault(VirtualFrame frame, PIBytesLike self, @SuppressWarnings("unused") PNone keepends) {
+            return doSplitlines(frame, self, false);
+        }
+
+        @Specialization(guards = "!isPNone(keepends)")
+        PList doSplitlinesDefault(VirtualFrame frame, PIBytesLike self, Object keepends,
+                        @Cached CastToJavaIntExactNode cast) {
+            return doSplitlines(frame, self, cast.execute(keepends) != 0);
+        }
+
+        @Specialization
+        PList doSplitlines(VirtualFrame frame, PIBytesLike self, boolean keepends) {
+            byte[] bytes = toBytesNode.execute(frame, self);
+            PList list = factory().createList();
+            int sliceStart = 0;
+            for (int i = 0; i < bytes.length; i++) {
+                if (bytes[i] == '\n' || bytes[i] == '\r') {
+                    int sliceEnd = i;
+                    if (bytes[i] == '\r' && i + 1 != bytes.length && bytes[i + 1] == '\n') {
+                        i++;
+                    }
+                    if (keepends) {
+                        sliceEnd = i + 1;
+                    }
+                    appendSlice(bytes, list, sliceStart, sliceEnd);
+                    sliceStart = i + 1;
+                }
+            }
+            // Process the remaining part if any
+            if (sliceStart != bytes.length) {
+                appendSlice(bytes, list, sliceStart, bytes.length);
+            }
+            return list;
+        }
+
+        private void appendSlice(byte[] bytes, PList list, int sliceStart, int sliceEnd) {
+            byte[] slice = new byte[sliceEnd - sliceStart];
+            PythonUtils.arraycopy(bytes, sliceStart, slice, 0, slice.length);
+            appendNode.execute(list, createElement(slice));
+        }
+
+        // Overriden in bytearray
+        protected PIBytesLike createElement(byte[] bytes) {
+            return factory().createBytes(bytes);
+        }
+    }
+
     // static bytes.maketrans()
     // static bytearray.maketrans()
     @Builtin(name = "maketrans", minNumOfPositionalArgs = 3, isClassmethod = true)