1 files changed, 43 insertions, 67 deletions

diff --git a/python/google/protobuf/pyext/extension_dict.cc b/python/google/protobuf/pyext/extension_dict.cc
index 3861c794..d83b57d5 100644
--- a/python/google/protobuf/pyext/extension_dict.cc
+++ b/python/google/protobuf/pyext/extension_dict.cc
@@ -62,22 +62,6 @@ static google::protobuf::Message* GetMessage(ExtensionDict* self) {
   }
 }
 
-CFieldDescriptor* InternalGetCDescriptorFromExtension(PyObject* extension) {
-  PyObject* cdescriptor = PyObject_GetAttrString(extension, "_cdescriptor");
-  if (cdescriptor == NULL) {
-    PyErr_SetString(PyExc_KeyError, "Unregistered extension.");
-    return NULL;
-  }
-  if (!PyObject_TypeCheck(cdescriptor, &CFieldDescriptor_Type)) {
-    PyErr_SetString(PyExc_TypeError, "Not a CFieldDescriptor");
-    Py_DECREF(cdescriptor);
-    return NULL;
-  }
-  CFieldDescriptor* descriptor =
-      reinterpret_cast<CFieldDescriptor*>(cdescriptor);
-  return descriptor;
-}
-
 PyObject* len(ExtensionDict* self) {
 #if PY_MAJOR_VERSION >= 3
   return PyLong_FromLong(PyDict_Size(self->values));
@@ -118,16 +102,15 @@ int ReleaseExtension(ExtensionDict* self,
 }
 
 PyObject* subscript(ExtensionDict* self, PyObject* key) {
-  CFieldDescriptor* cdescriptor = InternalGetCDescriptorFromExtension(
-      key);
-  if (cdescriptor == NULL) {
+  const google::protobuf::FieldDescriptor* descriptor =
+      cmessage::GetExtensionDescriptor(key);
+  if (descriptor == NULL) {
     return NULL;
   }
-  ScopedPyObjectPtr py_cdescriptor(reinterpret_cast<PyObject*>(cdescriptor));
-  const google::protobuf::FieldDescriptor* descriptor = cdescriptor->descriptor;
-  if (descriptor == NULL) {
+  if (!CheckFieldBelongsToMessage(descriptor, self->parent->message)) {
     return NULL;
   }
+
   if (descriptor->label() != FieldDescriptor::LABEL_REPEATED &&
       descriptor->cpp_type() != FieldDescriptor::CPPTYPE_MESSAGE) {
     return cmessage::InternalGetScalar(self->parent, descriptor);
@@ -142,7 +125,7 @@ PyObject* subscript(ExtensionDict* self, PyObject* key) {
   if (descriptor->label() != FieldDescriptor::LABEL_REPEATED &&
       descriptor->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
     PyObject* sub_message = cmessage::InternalGetSubMessage(
-        self->parent, cdescriptor);
+        self->parent, descriptor);
     if (sub_message == NULL) {
       return NULL;
     }
@@ -152,33 +135,21 @@ PyObject* subscript(ExtensionDict* self, PyObject* key) {
 
   if (descriptor->label() == FieldDescriptor::LABEL_REPEATED) {
     if (descriptor->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
-      // COPIED
-      PyObject* py_container = PyObject_CallObject(
-          reinterpret_cast<PyObject*>(&RepeatedCompositeContainer_Type),
-          NULL);
+      PyObject *message_class = cdescriptor_pool::GetMessageClass(
+          GetDescriptorPool(), descriptor->message_type());
+      if (message_class == NULL) {
+        return NULL;
+      }
+      PyObject* py_container = repeated_composite_container::NewContainer(
+          self->parent, descriptor, message_class);
       if (py_container == NULL) {
         return NULL;
       }
-      RepeatedCompositeContainer* container =
-          reinterpret_cast<RepeatedCompositeContainer*>(py_container);
-      PyObject* field = cdescriptor->descriptor_field;
-      PyObject* message_type = PyObject_GetAttrString(field, "message_type");
-      PyObject* concrete_class = PyObject_GetAttrString(message_type,
-                                                        "_concrete_class");
-      container->owner = self->owner;
-      container->parent = self->parent;
-      container->message = self->parent->message;
-      container->parent_field = cdescriptor;
-      container->subclass_init = concrete_class;
-      Py_DECREF(message_type);
       PyDict_SetItem(self->values, key, py_container);
       return py_container;
     } else {
-      // COPIED
-      ScopedPyObjectPtr init_args(PyTuple_Pack(2, self->parent, cdescriptor));
-      PyObject* py_container = PyObject_CallObject(
-          reinterpret_cast<PyObject*>(&RepeatedScalarContainer_Type),
-          init_args);
+      PyObject* py_container = repeated_scalar_container::NewContainer(
+          self->parent, descriptor);
       if (py_container == NULL) {
         return NULL;
       }
@@ -191,13 +162,15 @@ PyObject* subscript(ExtensionDict* self, PyObject* key) {
 }
 
 int ass_subscript(ExtensionDict* self, PyObject* key, PyObject* value) {
-  CFieldDescriptor* cdescriptor = InternalGetCDescriptorFromExtension(
-      key);
-  if (cdescriptor == NULL) {
+  const google::protobuf::FieldDescriptor* descriptor =
+      cmessage::GetExtensionDescriptor(key);
+  if (descriptor == NULL) {
+    return -1;
+  }
+  if (!CheckFieldBelongsToMessage(descriptor, self->parent->message)) {
     return -1;
   }
-  ScopedPyObjectPtr py_cdescriptor(reinterpret_cast<PyObject*>(cdescriptor));
-  const google::protobuf::FieldDescriptor* descriptor = cdescriptor->descriptor;
+
   if (descriptor->label() != FieldDescriptor::LABEL_OPTIONAL ||
       descriptor->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
     PyErr_SetString(PyExc_TypeError, "Extension is repeated and/or composite "
@@ -214,20 +187,18 @@ int ass_subscript(ExtensionDict* self, PyObject* key, PyObject* value) {
 }
 
 PyObject* ClearExtension(ExtensionDict* self, PyObject* extension) {
-  CFieldDescriptor* cdescriptor = InternalGetCDescriptorFromExtension(
-      extension);
-  if (cdescriptor == NULL) {
+  const google::protobuf::FieldDescriptor* descriptor =
+      cmessage::GetExtensionDescriptor(extension);
+  if (descriptor == NULL) {
     return NULL;
   }
-  ScopedPyObjectPtr py_cdescriptor(reinterpret_cast<PyObject*>(cdescriptor));
   PyObject* value = PyDict_GetItem(self->values, extension);
   if (value != NULL) {
-    if (ReleaseExtension(self, value, cdescriptor->descriptor) < 0) {
+    if (ReleaseExtension(self, value, descriptor) < 0) {
       return NULL;
     }
   }
-  if (cmessage::ClearFieldByDescriptor(self->parent,
-                                       cdescriptor->descriptor) == NULL) {
+  if (cmessage::ClearFieldByDescriptor(self->parent, descriptor) == NULL) {
     return NULL;
   }
   if (PyDict_DelItem(self->values, extension) < 0) {
@@ -237,14 +208,12 @@ PyObject* ClearExtension(ExtensionDict* self, PyObject* extension) {
 }
 
 PyObject* HasExtension(ExtensionDict* self, PyObject* extension) {
-  CFieldDescriptor* cdescriptor = InternalGetCDescriptorFromExtension(
-      extension);
-  if (cdescriptor == NULL) {
+  const google::protobuf::FieldDescriptor* descriptor =
+      cmessage::GetExtensionDescriptor(extension);
+  if (descriptor == NULL) {
     return NULL;
   }
-  ScopedPyObjectPtr py_cdescriptor(reinterpret_cast<PyObject*>(cdescriptor));
-  PyObject* result = cmessage::HasFieldByDescriptor(
-      self->parent, cdescriptor->descriptor);
+  PyObject* result = cmessage::HasFieldByDescriptor(self->parent, descriptor);
   return result;
 }
 
@@ -263,11 +232,18 @@ PyObject* _FindExtensionByName(ExtensionDict* self, PyObject* name) {
   }
 }
 
-int init(ExtensionDict* self, PyObject* args, PyObject* kwargs) {
-  self->parent = NULL;
-  self->message = NULL;
+ExtensionDict* NewExtensionDict(CMessage *parent) {
+  ExtensionDict* self = reinterpret_cast<ExtensionDict*>(
+      PyType_GenericAlloc(&ExtensionDict_Type, 0));
+  if (self == NULL) {
+    return NULL;
+  }
+
+  self->parent = parent;  // Store a borrowed reference.
+  self->message = parent->message;
+  self->owner = parent->owner;
   self->values = PyDict_New();
-  return 0;
+  return self;
 }
 
 void dealloc(ExtensionDict* self) {
@@ -330,7 +306,7 @@ PyTypeObject ExtensionDict_Type = {
   0,                                   //  tp_descr_get
   0,                                   //  tp_descr_set
   0,                                   //  tp_dictoffset
-  (initproc)extension_dict::init,      //  tp_init
+  0,                                   //  tp_init
 };
 
 }  // namespace python