1 files changed, 78 insertions, 17 deletions

diff --git a/python/google/protobuf/pyext/map_container.cc b/python/google/protobuf/pyext/map_container.cc
index 0987b898..318c2e7c 100644
--- a/python/google/protobuf/pyext/map_container.cc
+++ b/python/google/protobuf/pyext/map_container.cc
@@ -42,7 +42,9 @@
 #include <google/protobuf/map_field.h>
 #include <google/protobuf/map.h>
 #include <google/protobuf/message.h>
+#include <google/protobuf/pyext/message_factory.h>
 #include <google/protobuf/pyext/message.h>
+#include <google/protobuf/pyext/repeated_composite_container.h>
 #include <google/protobuf/pyext/scoped_pyobject_ptr.h>
 
 #if PY_MAJOR_VERSION >= 3
@@ -328,6 +330,15 @@ PyObject* Clear(PyObject* _self) {
   Py_RETURN_NONE;
 }
 
+PyObject* GetEntryClass(PyObject* _self) {
+  MapContainer* self = GetMap(_self);
+  CMessageClass* message_class = message_factory::GetMessageClass(
+      cmessage::GetFactoryForMessage(self->parent),
+      self->parent_field_descriptor->message_type());
+  Py_XINCREF(message_class);
+  return reinterpret_cast<PyObject*>(message_class);
+}
+
 PyObject* MapReflectionFriend::Contains(PyObject* _self, PyObject* key) {
   MapContainer* self = GetMap(_self);
 
@@ -400,12 +411,7 @@ PyObject *NewScalarMapContainer(
     return NULL;
   }
 
-#if PY_MAJOR_VERSION >= 3
-  ScopedPyObjectPtr obj(PyType_GenericAlloc(
-        reinterpret_cast<PyTypeObject *>(ScalarMapContainer_Type), 0));
-#else
-  ScopedPyObjectPtr obj(PyType_GenericAlloc(&ScalarMapContainer_Type, 0));
-#endif
+  ScopedPyObjectPtr obj(PyType_GenericAlloc(ScalarMapContainer_Type, 0));
   if (obj.get() == NULL) {
     return PyErr_Format(PyExc_RuntimeError,
                         "Could not allocate new container.");
@@ -527,6 +533,8 @@ static PyMethodDef ScalarMapMethods[] = {
     "Removes all elements from the map." },
   { "get", ScalarMapGet, METH_VARARGS,
     "Gets the value for the given key if present, or otherwise a default" },
+  { "GetEntryClass", (PyCFunction)GetEntryClass, METH_NOARGS,
+    "Return the class used to build Entries of (key, value) pairs." },
   /*
   { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
     "Makes a deep copy of the class." },
@@ -536,6 +544,7 @@ static PyMethodDef ScalarMapMethods[] = {
   {NULL, NULL},
 };
 
+PyTypeObject *ScalarMapContainer_Type;
 #if PY_MAJOR_VERSION >= 3
   static PyType_Slot ScalarMapContainer_Type_slots[] = {
       {Py_tp_dealloc, (void *)ScalarMapDealloc},
@@ -554,7 +563,6 @@ static PyMethodDef ScalarMapMethods[] = {
       Py_TPFLAGS_DEFAULT,
       ScalarMapContainer_Type_slots
   };
-  PyObject *ScalarMapContainer_Type;
 #else
   static PyMappingMethods ScalarMapMappingMethods = {
     MapReflectionFriend::Length,             // mp_length
@@ -562,7 +570,7 @@ static PyMethodDef ScalarMapMethods[] = {
     MapReflectionFriend::ScalarMapSetItem,   // mp_ass_subscript
   };
 
-  PyTypeObject ScalarMapContainer_Type = {
+  PyTypeObject _ScalarMapContainer_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     FULL_MODULE_NAME ".ScalarMapContainer",  //  tp_name
     sizeof(MapContainer),                //  tp_basicsize
@@ -643,12 +651,7 @@ PyObject* NewMessageMapContainer(
     return NULL;
   }
 
-#if PY_MAJOR_VERSION >= 3
-  PyObject* obj = PyType_GenericAlloc(
-        reinterpret_cast<PyTypeObject *>(MessageMapContainer_Type), 0);
-#else
-  PyObject* obj = PyType_GenericAlloc(&MessageMapContainer_Type, 0);
-#endif
+  PyObject* obj = PyType_GenericAlloc(MessageMapContainer_Type, 0);
   if (obj == NULL) {
     return PyErr_Format(PyExc_RuntimeError,
                         "Could not allocate new container.");
@@ -780,6 +783,8 @@ static PyMethodDef MessageMapMethods[] = {
     "Gets the value for the given key if present, or otherwise a default" },
   { "get_or_create", MapReflectionFriend::MessageMapGetItem, METH_O,
     "Alias for getitem, useful to make explicit that the map is mutated." },
+  { "GetEntryClass", (PyCFunction)GetEntryClass, METH_NOARGS,
+    "Return the class used to build Entries of (key, value) pairs." },
   /*
   { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
     "Makes a deep copy of the class." },
@@ -789,6 +794,7 @@ static PyMethodDef MessageMapMethods[] = {
   {NULL, NULL},
 };
 
+PyTypeObject *MessageMapContainer_Type;
 #if PY_MAJOR_VERSION >= 3
   static PyType_Slot MessageMapContainer_Type_slots[] = {
       {Py_tp_dealloc, (void *)MessageMapDealloc},
@@ -807,8 +813,6 @@ static PyMethodDef MessageMapMethods[] = {
       Py_TPFLAGS_DEFAULT,
       MessageMapContainer_Type_slots
   };
-
-  PyObject *MessageMapContainer_Type;
 #else
   static PyMappingMethods MessageMapMappingMethods = {
     MapReflectionFriend::Length,              // mp_length
@@ -816,7 +820,7 @@ static PyMethodDef MessageMapMethods[] = {
     MapReflectionFriend::MessageMapSetItem,   // mp_ass_subscript
   };
 
-  PyTypeObject MessageMapContainer_Type = {
+  PyTypeObject _MessageMapContainer_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     FULL_MODULE_NAME ".MessageMapContainer",  //  tp_name
     sizeof(MessageMapContainer),         //  tp_basicsize
@@ -965,6 +969,63 @@ PyTypeObject MapIterator_Type = {
   0,                                   //  tp_init
 };
 
+bool InitMapContainers() {
+  // ScalarMapContainer_Type derives from our MutableMapping type.
+  ScopedPyObjectPtr containers(PyImport_ImportModule(
+      "google.protobuf.internal.containers"));
+  if (containers == NULL) {
+    return false;
+  }
+
+  ScopedPyObjectPtr mutable_mapping(
+      PyObject_GetAttrString(containers.get(), "MutableMapping"));
+  if (mutable_mapping == NULL) {
+    return false;
+  }
+
+  if (!PyObject_TypeCheck(mutable_mapping.get(), &PyType_Type)) {
+    return false;
+  }
+
+  Py_INCREF(mutable_mapping.get());
+#if PY_MAJOR_VERSION >= 3
+  PyObject* bases = PyTuple_New(1);
+  PyTuple_SET_ITEM(bases, 0, mutable_mapping.get());
+
+  ScalarMapContainer_Type = reinterpret_cast<PyTypeObject*>(
+      PyType_FromSpecWithBases(&ScalarMapContainer_Type_spec, bases));
+#else
+  _ScalarMapContainer_Type.tp_base =
+      reinterpret_cast<PyTypeObject*>(mutable_mapping.get());
+
+  if (PyType_Ready(&_ScalarMapContainer_Type) < 0) {
+    return false;
+  }
+
+  ScalarMapContainer_Type = &_ScalarMapContainer_Type;
+#endif
+
+  if (PyType_Ready(&MapIterator_Type) < 0) {
+    return false;
+  }
+
+#if PY_MAJOR_VERSION >= 3
+  MessageMapContainer_Type = reinterpret_cast<PyTypeObject*>(
+      PyType_FromSpecWithBases(&MessageMapContainer_Type_spec, bases));
+#else
+  Py_INCREF(mutable_mapping.get());
+  _MessageMapContainer_Type.tp_base =
+      reinterpret_cast<PyTypeObject*>(mutable_mapping.get());
+
+  if (PyType_Ready(&_MessageMapContainer_Type) < 0) {
+    return false;
+  }
+
+  MessageMapContainer_Type = &_MessageMapContainer_Type;
+#endif
+  return true;
+}
+
 }  // namespace python
 }  // namespace protobuf
 }  // namespace google