diff --git a/python/google/protobuf/internal/message_test.py b/python/google/protobuf/internal/message_test.py index ebb51552c0008..dfd948c6ce2a4 100755 --- a/python/google/protobuf/internal/message_test.py +++ b/python/google/protobuf/internal/message_test.py @@ -1058,6 +1058,13 @@ def testExtendStringWithNothing(self, message_module): m.repeated_string.extend(empty_value) self.assertSequenceEqual([], m.repeated_string) + def testExtendWithNoLen(self, message_module): + """ Test extending repeated fields with iterables but no len""" + m = message_module.TestAllTypes() + self.assertSequenceEqual([], m.repeated_int32) + m.repeated_int32.extend(i for i in range(2)) + self.assertSequenceEqual([0, 1], m.repeated_int32) + def testExtendInt32WithPythonList(self, message_module): """Test extending repeated int32 fields with python lists.""" m = message_module.TestAllTypes() @@ -1068,6 +1075,8 @@ def testExtendInt32WithPythonList(self, message_module): self.assertSequenceEqual([0, 1, 2], m.repeated_int32) m.repeated_int32.extend([3, 4]) self.assertSequenceEqual([0, 1, 2, 3, 4], m.repeated_int32) + with self.assertRaises(TypeError): + m.repeated_int32.extend([5, 6, 'hi', 7]) def testExtendFloatWithPythonList(self, message_module): """Test extending repeated float fields with python lists.""" diff --git a/python/repeated.c b/python/repeated.c index c881ecebe8b75..c0e39c079e458 100644 --- a/python/repeated.c +++ b/python/repeated.c @@ -181,15 +181,32 @@ PyObject* PyUpb_RepeatedContainer_Extend(PyObject* _self, PyObject* value) { bool submsg = upb_FieldDef_IsSubMessage(f); PyObject* e; - while ((e = PyIter_Next(it))) { - PyObject* ret; - if (submsg) { - ret = PyUpb_RepeatedCompositeContainer_Append(_self, e); + if (submsg) { + while ((e = PyIter_Next(it))) { + PyObject* ret = PyUpb_RepeatedCompositeContainer_Append(_self, e); + Py_XDECREF(ret); + Py_DECREF(e); + } + } else { + upb_Arena* arena = PyUpb_Arena_Get(self->arena); + Py_ssize_t size = PyObject_Size(value); + if (size < 0) { + // Some iterables may not have len. Size() will return -1 and + // set an error in such cases. + PyErr_Clear(); } else { - ret = PyUpb_RepeatedScalarContainer_Append(_self, e); + upb_Array_Reserve(arr, start_size + size, arena); + } + while ((e = PyIter_Next(it))) { + upb_MessageValue msgval; + if (!PyUpb_PyToUpb(e, f, &msgval, arena)) { + assert(PyErr_Occurred()); + Py_DECREF(e); + break; + } + upb_Array_Append(arr, msgval, arena); + Py_DECREF(e); } - Py_XDECREF(ret); - Py_DECREF(e); } Py_DECREF(it);