Fix python proto_api to work with messages that required fields are not set.

anandolee · copybara-github · commit f63ed55d4b85 · 2025-05-02T13:49:04.000-07:00
Fix python proto_api to work with large messages for cpp extension (upb and
pure python can not work with large messages which are larger than 2GB)

PiperOrigin-RevId: 754120424
diff --git a/python/google/protobuf/proto_api.cc b/python/google/protobuf/proto_api.cc
@@ -2,6 +2,7 @@
 
 #include <Python.h>
 
+#include <climits>
 #include <memory>
 #include <string>
 
@@ -38,7 +39,7 @@ PythonMessageMutator::~PythonMessageMutator() {
   // check.
   if (!PyErr_Occurred() && owned_msg_ != nullptr) {
     std::string wire;
-    message_->SerializeToString(&wire);
+    message_->SerializePartialToString(&wire);
     PyObject* py_wire = PyBytes_FromStringAndSize(
         wire.data(), static_cast<Py_ssize_t>(wire.size()));
     PyObject* parse =
@@ -81,8 +82,14 @@ bool PythonConstMessagePointer::NotChanged() {
     return false;
   }
 
+  // Skip the check if too large. Parse won't work
+  // for messages larger than 2 GB.
+  if (message_->ByteSizeLong() > INT_MAX) {
+    return true;
+  }
+
   PyObject* py_serialized_pb(
-      PyObject_CallMethod(py_msg_, "SerializeToString", nullptr));
+      PyObject_CallMethod(py_msg_, "SerializePartialToString", nullptr));
   if (py_serialized_pb == nullptr) {
     PyErr_Format(PyExc_ValueError, "Fail to serialize py_msg");
     return false;
@@ -99,19 +106,19 @@ bool PythonConstMessagePointer::NotChanged() {
   // serialize result may still diff between languages. So parse to
   // another c++ message for compare.
   std::unique_ptr<google::protobuf::Message> parsed_msg(owned_msg_->New());
-  parsed_msg->ParseFromArray(data, static_cast<int>(len));
+  parsed_msg->ParsePartialFromArray(data, static_cast<int>(len));
   std::string wire_other;
   google::protobuf::io::StringOutputStream stream_other(&wire_other);
   google::protobuf::io::CodedOutputStream output_other(&stream_other);
   output_other.SetSerializationDeterministic(true);
-  parsed_msg->SerializeToCodedStream(&output_other);
+  parsed_msg->SerializePartialToCodedStream(&output_other);
   output_other.Trim();
 
   std::string wire;
   google::protobuf::io::StringOutputStream stream(&wire);
   google::protobuf::io::CodedOutputStream output(&stream);
   output.SetSerializationDeterministic(true);
-  owned_msg_->SerializeToCodedStream(&output);
+  owned_msg_->SerializePartialToCodedStream(&output);
   output.Trim();
 
   if (wire == wire_other) {
diff --git a/python/google/protobuf/pyext/message_module.cc b/python/google/protobuf/pyext/message_module.cc
@@ -85,7 +85,7 @@ class ProtoAPIDescriptorDatabase : public google::protobuf::DescriptorDatabase {
       return false;
     }
 
-    bool ok = output->ParseFromArray(
+    bool ok = output->ParsePartialFromArray(
         reinterpret_cast<uint8_t*>(PyBytes_AS_STRING(pyfile_serialized)),
         PyBytes_GET_SIZE(pyfile_serialized));
     if (!ok) {
@@ -207,8 +207,8 @@ absl::StatusOr<google::protobuf::Message*> CreateNewMessage(PyObject* py_msg) {
 bool CopyToOwnedMsg(google::protobuf::Message** copy, const google::protobuf::Message& message) {
   *copy = message.New();
   std::string wire;
-  message.SerializeToString(&wire);
-  (*copy)->ParseFromArray(wire.data(), wire.size());
+  message.SerializePartialToString(&wire);
+  (*copy)->ParsePartialFromArray(wire.data(), wire.size());
   return true;
 }
 
@@ -245,7 +245,7 @@ struct ApiImplementation : google::protobuf::python::PyProto_API {
     auto msg = CreateNewMessage(py_msg);
     RETURN_IF_ERROR(msg.status());
     PyObject* serialized_pb(
-        PyObject_CallMethod(py_msg, "SerializeToString", nullptr));
+        PyObject_CallMethod(py_msg, "SerializePartialToString", nullptr));
     if (serialized_pb == nullptr) {
       return absl::InternalError("Fail to serialize py_msg");
     }
@@ -256,7 +256,7 @@ struct ApiImplementation : google::protobuf::python::PyProto_API {
       return absl::InternalError(
           "Fail to get bytes from py_msg serialized data");
     }
-    if (!(*msg)->ParseFromArray(data, len)) {
+    if (!(*msg)->ParsePartialFromArray(data, len)) {
       Py_DECREF(serialized_pb);
       return absl::InternalError(
           "Couldn't parse py_message to google::protobuf::Message*!");
