[mypyc] Support bytes concat

python · JukkaL · Aug 6, 2021 · Aug 1, 2021 · Aug 3, 2021 · Aug 3, 2021
commit e871d1b69d0077640fcc41935bdf75fb4ab8ff1d
diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
@@ -400,6 +400,8 @@ Py_ssize_t CPyStr_Size_size_t(PyObject *str);
 // Bytes operations
 
 
+PyObject *CPyBytes_Concat(PyObject *a, PyObject *b);
+
 
 // Set operations
 

diff --git a/mypyc/lib-rt/bytes_ops.c b/mypyc/lib-rt/bytes_ops.c
@@ -4,3 +4,12 @@
 
 #include <Python.h>
 #include "CPy.h"
+
+PyObject *CPyBytes_Concat(PyObject *a, PyObject *b) {
+    if (PyByteArray_Check(a) && PyByteArray_Check(b)) {
+        return PyByteArray_Concat(a, b);
+    } else {
+        PyBytes_Concat(&a, b);
+        return a;
+    }
+}
diff --git a/mypyc/primitives/bytes_ops.py b/mypyc/primitives/bytes_ops.py
@@ -5,8 +5,7 @@
     object_rprimitive, bytes_rprimitive, list_rprimitive, dict_rprimitive,
     str_rprimitive, RUnion
 )
-from mypyc.primitives.registry import load_address_op, function_op
-
+from mypyc.primitives.registry import load_address_op, function_op, binary_op
 
 # Get the 'bytes' type object.
 load_address_op(
@@ -29,3 +28,13 @@
     return_type=bytes_rprimitive,
     c_function_name='PyByteArray_FromObject',
     error_kind=ERR_MAGIC)
+
+# bytes + bytes
+# bytearray + bytearray
+binary_op(
+    name='+',
+    arg_types=[bytes_rprimitive, bytes_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name='CPyBytes_Concat',
+    error_kind=ERR_MAGIC,
+    steals=[True, False])
diff --git a/mypyc/test-data/irbuild-bytes.test b/mypyc/test-data/irbuild-bytes.test
@@ -62,3 +62,12 @@ L0:
     c = r6
     return 1
 
+[case testBytesOps]
+def f1(a: bytes, b: bytes) -> bytes:
+    return a + b
+[out]
+def f1(a, b):
+    a, b, r0 :: bytes
+L0:
+    r0 = CPyBytes_Concat(a, b)
+    return r0
diff --git a/mypyc/test-data/run-bytes.test b/mypyc/test-data/run-bytes.test
@@ -58,6 +58,24 @@ def test_concat() -> None:
     b1 = b'123' + bytes()
     b2 = b'456' + bytes()
     assert b1 + b2 == b'123456'
+    b3 = b1 + b2
+    b3 = b3 + b1
+    assert b3 == b'123456123'
+    assert b1 == b'123'
+    assert b2 == b'456'
+    brr1: bytes = bytearray(3)
+    brr2: bytes = bytearray(range(5))
+    b4 = b1 + brr1
+    assert b4 == b'123\x00\x00\x00'
+    brr3 = brr1 + brr2
+    assert brr3 == bytearray(b'\x00\x00\x00\x00\x01\x02\x03\x04')
+    assert len(brr3) == 8
+    brr3 = brr3 + bytearray([10])
+    assert brr3 == bytearray(b'\x00\x00\x00\x00\x01\x02\x03\x04\n')
+    b5 = brr2 + b2
+    assert b5 == bytearray(b'\x00\x01\x02\x03\x04456')
+    b5 = b2 + brr2
+    assert b5 == b'456\x00\x01\x02\x03\x04'
 
 def test_join() -> None:
     seq = (b'1', b'"', b'\xf0')
Original file line number	Diff line number	Diff line change
Expand Up		@@ -400,6 +400,8 @@ Py_ssize_t CPyStr_Size_size_t(PyObject *str);
		// Bytes operations


		PyObject CPyBytes_Concat(PyObject a, PyObject *b);


		// Set operations

Expand Down