[GR-13242] Bytearray += operation has to return the same object.

PullRequest: graalpython/348

Author: ppisl

graalpython/com.oracle.graal.python.test/src/tests/test_bytes.py

@@ -610,6 +610,14 @@ class BytesSplitTest(BaseTestSplit, unittest.TestCase):
 class ByteArraySplitTest(BaseTestSplit, unittest.TestCase):
     type2test = bytearray
 
+
+def test_eq_add_bytearray():
+    b1 = bytearray(b'')
+    b2 = b1
+    b1 += bytearray(b'Hello')
+    assert b1 is b2
+
+
 def test_add_mv_to_bytes():
     b = b'hello '
     mv = memoryview(b'world')

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

@@ -28,6 +28,7 @@
 
 import static com.oracle.graal.python.builtins.objects.slice.PSlice.MISSING_INDEX;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__ADD__;
+import static com.oracle.graal.python.nodes.SpecialMethodNames.__IADD__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__BOOL__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__DELITEM__;
 import static com.oracle.graal.python.nodes.SpecialMethodNames.__EQ__;
@@ -265,6 +266,45 @@ public Object add(Object self, Object other) {
         }
     }
 
+    @Builtin(name = __IADD__, fixedNumOfPositionalArgs = 2)
+    @GenerateNodeFactory
+    public abstract static class IAddNode extends PythonBinaryBuiltinNode {
+        @Specialization
+        public PByteArray add(PByteArray self, PIBytesLike other,
+                        @Cached("create()") SequenceStorageNodes.ConcatNode concatNode) {
+            SequenceStorage res = concatNode.execute(self.getSequenceStorage(), other.getSequenceStorage());
+            updateSequenceStorage(self, res);
+            return self;
+        }
+
+        @Specialization
+        public PByteArray add(PByteArray self, PMemoryView other,
+                        @Cached("create(TOBYTES)") LookupAndCallUnaryNode toBytesNode,
+                        @Cached("createBinaryProfile()") ConditionProfile isBytesProfile,
+                        @Cached("create()") SequenceStorageNodes.ConcatNode concatNode) {
+
+            Object bytesObj = toBytesNode.executeObject(other);
+            if (isBytesProfile.profile(bytesObj instanceof PBytes)) {
+                SequenceStorage res = concatNode.execute(self.getSequenceStorage(), ((PBytes) bytesObj).getSequenceStorage());
+                updateSequenceStorage(self, res);
+                return self;
+            }
+            throw raise(SystemError, "could not get bytes of memoryview");
+        }
+
+        @SuppressWarnings("unused")
+        @Fallback
+        public Object add(Object self, Object other) {
+            throw raise(TypeError, "can't concat bytearray to %p", other);
+        }
+
+        private static void updateSequenceStorage(PByteArray array, SequenceStorage s) {
+            if (array.getSequenceStorage() != s) {
+                array.setSequenceStorage(s);
+            }
+        }
+    }
+
     @Builtin(name = __MUL__, fixedNumOfPositionalArgs = 2)
     @GenerateNodeFactory
     public abstract static class MulNode extends PythonBuiltinNode {