From ea4cb79f669e69342d7ced4d0255050325df41e3 Mon Sep 17 00:00:00 2001
From: Eric Salo <salo@google.com>
Date: Wed, 22 Mar 2023 07:25:40 -0700
Subject: [PATCH] fix Python bug with required fields

https://github.com/protocolbuffers/upb/issues/1220

There were two bugs here: Python was incorrectly mandating that a required
field be set during assignment, and it was also incorrectly assuming a non-NULL
return from an internal function call.

PiperOrigin-RevId: 518561818
---
 python/message.c | 22 ++++++++++++++++++----
 python/message.h |  2 ++
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/python/message.c b/python/message.c
index 6bd3941646..753cc7453e 100644
--- a/python/message.c
+++ b/python/message.c
@@ -447,6 +447,8 @@ err:
   return ok;
 }
 
+static PyObject* PyUpb_Message_MergePartialFrom(PyObject*, PyObject*);
+
 static bool PyUpb_Message_InitMessageAttribute(PyObject* _self, PyObject* name,
                                                PyObject* value) {
   PyObject* submsg = PyUpb_Message_GetAttr(_self, name);
@@ -454,9 +456,9 @@ static bool PyUpb_Message_InitMessageAttribute(PyObject* _self, PyObject* name,
   assert(!PyErr_Occurred());
   bool ok;
   if (PyUpb_Message_TryCheck(value)) {
-    PyObject* tmp = PyUpb_Message_MergeFrom(submsg, value);
+    PyObject* tmp = PyUpb_Message_MergePartialFrom(submsg, value);
     ok = tmp != NULL;
-    Py_DECREF(tmp);
+    Py_XDECREF(tmp);
   } else if (PyDict_Check(value)) {
     assert(!PyErr_Occurred());
     ok = PyUpb_Message_InitAttributes(submsg, NULL, value) >= 0;
@@ -1184,7 +1186,8 @@ err:
   return NULL;
 }
 
-PyObject* PyUpb_Message_MergeFrom(PyObject* self, PyObject* arg) {
+static PyObject* PyUpb_Message_MergeInternal(PyObject* self, PyObject* arg,
+                                             bool check_required) {
   if (self->ob_type != arg->ob_type) {
     PyErr_Format(PyExc_TypeError,
                  "Parameter to MergeFrom() must be instance of same class: "
@@ -1194,7 +1197,10 @@ PyObject* PyUpb_Message_MergeFrom(PyObject* self, PyObject* arg) {
   }
   // OPT: exit if src is empty.
   PyObject* subargs = PyTuple_New(0);
-  PyObject* serialized = PyUpb_Message_SerializeToString(arg, subargs, NULL);
+  PyObject* serialized =
+      check_required
+          ? PyUpb_Message_SerializeToString(arg, subargs, NULL)
+          : PyUpb_Message_SerializePartialToString(arg, subargs, NULL);
   Py_DECREF(subargs);
   if (!serialized) return NULL;
   PyObject* ret = PyUpb_Message_MergeFromString(self, serialized);
@@ -1203,6 +1209,14 @@ PyObject* PyUpb_Message_MergeFrom(PyObject* self, PyObject* arg) {
   Py_RETURN_NONE;
 }
 
+PyObject* PyUpb_Message_MergeFrom(PyObject* self, PyObject* arg) {
+  return PyUpb_Message_MergeInternal(self, arg, true);
+}
+
+static PyObject* PyUpb_Message_MergePartialFrom(PyObject* self, PyObject* arg) {
+  return PyUpb_Message_MergeInternal(self, arg, false);
+}
+
 static PyObject* PyUpb_Message_SetInParent(PyObject* _self, PyObject* arg) {
   PyUpb_Message* self = (void*)_self;
   PyUpb_Message_EnsureReified(self);
diff --git a/python/message.h b/python/message.h
index f10f15f02a..2bb075ba30 100644
--- a/python/message.h
+++ b/python/message.h
@@ -64,6 +64,8 @@ PyObject* PyUpb_Message_MergeFrom(PyObject* self, PyObject* arg);
 PyObject* PyUpb_Message_MergeFromString(PyObject* self, PyObject* arg);
 PyObject* PyUpb_Message_SerializeToString(PyObject* self, PyObject* args,
                                           PyObject* kwargs);
+PyObject* PyUpb_Message_SerializePartialToString(PyObject* self, PyObject* args,
+                                                 PyObject* kwargs);
 
 // Sets fields of the message according to the attribuges in `kwargs`.
 int PyUpb_Message_InitAttributes(PyObject* _self, PyObject* args,