diff options
Diffstat (limited to 'python/google/protobuf/pyext/message.cc')
-rw-r--r-- | python/google/protobuf/pyext/message.cc | 62 |
1 files changed, 48 insertions, 14 deletions
diff --git a/python/google/protobuf/pyext/message.cc b/python/google/protobuf/pyext/message.cc index 49113c7c..702c5d03 100644 --- a/python/google/protobuf/pyext/message.cc +++ b/python/google/protobuf/pyext/message.cc @@ -52,6 +52,7 @@ #include <google/protobuf/stubs/common.h> #include <google/protobuf/stubs/logging.h> #include <google/protobuf/io/coded_stream.h> +#include <google/protobuf/io/zero_copy_stream_impl_lite.h> #include <google/protobuf/util/message_differencer.h> #include <google/protobuf/descriptor.h> #include <google/protobuf/message.h> @@ -1810,8 +1811,25 @@ static string GetMessageName(CMessage* self) { } } -static PyObject* SerializeToString(CMessage* self, PyObject* args) { - if (!self->message->IsInitialized()) { +static PyObject* InternalSerializeToString( + CMessage* self, PyObject* args, PyObject* kwargs, + bool require_initialized) { + // Parse the "deterministic" kwarg; defaults to False. + static char* kwlist[] = { "deterministic", 0 }; + PyObject* deterministic_obj = Py_None; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O", kwlist, + &deterministic_obj)) { + return NULL; + } + // Preemptively convert to a bool first, so we don't need to back out of + // allocating memory if this raises an exception. + // NOTE: This is unused later if deterministic == Py_None, but that's fine. + int deterministic = PyObject_IsTrue(deterministic_obj); + if (deterministic < 0) { + return NULL; + } + + if (require_initialized && !self->message->IsInitialized()) { ScopedPyObjectPtr errors(FindInitializationErrors(self)); if (errors == NULL) { return NULL; @@ -1849,24 +1867,36 @@ static PyObject* SerializeToString(CMessage* self, PyObject* args) { GetMessageName(self).c_str(), PyString_AsString(joined.get())); return NULL; } - int size = self->message->ByteSize(); - if (size <= 0) { + + // Ok, arguments parsed and errors checked, now encode to a string + const size_t size = self->message->ByteSizeLong(); + if (size == 0) { return PyBytes_FromString(""); } PyObject* result = PyBytes_FromStringAndSize(NULL, size); if (result == NULL) { return NULL; } - char* buffer = PyBytes_AS_STRING(result); - self->message->SerializeWithCachedSizesToArray( - reinterpret_cast<uint8*>(buffer)); + io::ArrayOutputStream out(PyBytes_AS_STRING(result), size); + io::CodedOutputStream coded_out(&out); + if (deterministic_obj != Py_None) { + coded_out.SetSerializationDeterministic(deterministic); + } + self->message->SerializeWithCachedSizes(&coded_out); + GOOGLE_CHECK(!coded_out.HadError()); return result; } -static PyObject* SerializePartialToString(CMessage* self) { - string contents; - self->message->SerializePartialToString(&contents); - return PyBytes_FromStringAndSize(contents.c_str(), contents.size()); +static PyObject* SerializeToString( + CMessage* self, PyObject* args, PyObject* kwargs) { + return InternalSerializeToString(self, args, kwargs, + /*require_initialized=*/true); +} + +static PyObject* SerializePartialToString( + CMessage* self, PyObject* args, PyObject* kwargs) { + return InternalSerializeToString(self, args, kwargs, + /*require_initialized=*/false); } // Formats proto fields for ascii dumps using python formatting functions where @@ -2537,7 +2567,10 @@ PyObject* Reduce(CMessage* self) { if (state == NULL) { return NULL; } - ScopedPyObjectPtr serialized(SerializePartialToString(self)); + string contents; + self->message->SerializePartialToString(&contents); + ScopedPyObjectPtr serialized( + PyBytes_FromStringAndSize(contents.c_str(), contents.size())); if (serialized == NULL) { return NULL; } @@ -2658,9 +2691,10 @@ static PyMethodDef Methods[] = { { "RegisterExtension", (PyCFunction)RegisterExtension, METH_O | METH_CLASS, "Registers an extension with the current message." }, { "SerializePartialToString", (PyCFunction)SerializePartialToString, - METH_NOARGS, + METH_VARARGS | METH_KEYWORDS, "Serializes the message to a string, even if it isn't initialized." }, - { "SerializeToString", (PyCFunction)SerializeToString, METH_NOARGS, + { "SerializeToString", (PyCFunction)SerializeToString, + METH_VARARGS | METH_KEYWORDS, "Serializes the message to a string, only for initialized messages." }, { "SetInParent", (PyCFunction)SetInParent, METH_NOARGS, "Sets the has bit of the given field in its parent message." }, |