make bytearray work (again)

protocolbuffers · May 30, 2024 · 6249e62 · 6249e62
1 parent 092cc15
commit 6249e62
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/python/google/protobuf/internal/message_test.py b/python/google/protobuf/internal/message_test.py
@@ -82,6 +82,19 @@ def testGoldenMessage(self, message_module):
     golden_copy = copy.deepcopy(golden_message)
     self.assertEqual(golden_data, golden_copy.SerializeToString())
 
+  def testGoldenMessageBytearray(self, message_module):
+    # bytearray was broken, test that it works again
+    if message_module is unittest_pb2:
+      golden_data = test_util.GoldenFileData('golden_message_oneof_implemented')
+    else:
+      golden_data = test_util.GoldenFileData('golden_message_proto3')
+
+    golden_message = message_module.TestAllTypes()
+    golden_message.ParseFromString(bytearray(golden_data))
+    if message_module is unittest_pb2:
+      test_util.ExpectAllFieldsSet(self, golden_message)
+    self.assertEqual(golden_data, golden_message.SerializeToString())
+
   def testGoldenPackedMessage(self, message_module):
     golden_data = test_util.GoldenFileData('golden_packed_fields_message')
     golden_message = message_module.TestPackedTypes()

diff --git a/python/message.c b/python/message.c
@@ -1286,6 +1286,9 @@ PyObject* PyUpb_Message_MergeFromString(PyObject* _self, PyObject* arg) {
     int err = PyBytes_AsStringAndSize(bytes, &buf, &size);
     (void)err;
     assert(err >= 0);
+  } else if (PyByteArray_Check(arg)) {
+    buf = PyByteArray_AS_STRING(arg);
+    size = PyByteArray_GET_SIZE(arg);
   } else if (PyBytes_AsStringAndSize(arg, &buf, &size) < 0) {
     return NULL;
   }