From 6c3ff5aa9eba8fd78f8b52afe98df11c1fb85ad2 Mon Sep 17 00:00:00 2001
From: Jie Luo <jieluo@google.com>
Date: Fri, 3 Jan 2025 18:20:05 -0800
Subject: [PATCH] Performance optimization for upb python extend repeated
 scalars.

Before:
_bench_append, 100000, 241.9
_bench_extend, 100000, 112.3
_bench_assign, 100000, 83.3
_bench_pybind11, 100000, 16.2

After:
_bench_append, 100000, 224.3
_bench_extend, 100000, 71.2
_bench_assign, 100000, 81.1
_bench_pybind11, 100000, 14.7
PiperOrigin-RevId: 711909495
---
 .../google/protobuf/internal/message_test.py  |  9 ++++++
 python/repeated.c                             | 31 ++++++++++++++-----
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/python/google/protobuf/internal/message_test.py b/python/google/protobuf/internal/message_test.py
index ebb51552c0..dfd948c6ce 100755
--- a/python/google/protobuf/internal/message_test.py
+++ b/python/google/protobuf/internal/message_test.py
@@ -1058,6 +1058,13 @@ class MessageTest(unittest.TestCase):
       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 @@ class MessageTest(unittest.TestCase):
     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 c881ecebe8..c0e39c079e 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);