Added beginnings of support for descriptor DB.

python/descriptor_pool.c

@@ -121,30 +121,53 @@ static PyObject* PyUpb_DescriptorPool_New(PyTypeObject* type, PyObject* args,
   return PyUpb_DescriptorPool_DoCreate(type, db);
 }
 
-/*
- * PyUpb_DescriptorPool_AddSerializedFile()
- *
- * Implements:
- *   DescriptorPool.AddSerializedFile(self, serialized_file_descriptor)
- *
- * Adds the given serialized FileDescriptorProto to the pool.
- */
-static PyObject* PyUpb_DescriptorPool_AddSerializedFile(
-    PyObject* _self, PyObject* serialized_pb) {
-  PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
-  char* buf;
-  Py_ssize_t size;
-  upb_arena* arena = upb_arena_new();
-  if (!arena) PYUPB_RETURN_OOM;
-  PyObject* result = NULL;
+static PyObject* PyUpb_DescriptorPool_DoAdd(PyObject* _self,
+                                            PyObject* file_desc);
+
+static bool PyUpb_DescriptorPool_TryLoadFileProto(PyUpb_DescriptorPool* self,
+                                                  PyObject* proto) {
+  if (proto == NULL) {
+    if (PyErr_ExceptionMatches(PyExc_KeyError)) {
+      // Expected error: item was simply not found.
+      PyErr_Clear();
+      return true;  // We didn't accomplish our goal, but we didn't error out.
+    }
+    return false;
+  }
+  if (proto == Py_None) return true;
+  return PyUpb_DescriptorPool_DoAdd((PyObject*)self, proto) != NULL;
+}
 
+static bool PyUpb_DescriptorPool_TryLoadSymbol(PyUpb_DescriptorPool* self,
+                                               PyObject* sym) {
+  if (!self->db) return false;
+  PyObject* file_proto =
+      PyObject_CallMethod(self->db, "FindFileContainingSymbol", "O", sym);
+  bool ret = PyUpb_DescriptorPool_TryLoadFileProto(self, file_proto);
+  Py_XDECREF(file_proto);
+  return ret;
+}
+
+bool PyUpb_DescriptorPool_CheckNoDatabase(PyObject* _self) {
+  PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
   if (self->db) {
     PyErr_SetString(
         PyExc_ValueError,
         "Cannot call Add on a DescriptorPool that uses a DescriptorDatabase. "
         "Add your file to the underlying database.");
-    goto done;
+    return false;
   }
+  return true;
+}
+
+static PyObject* PyUpb_DescriptorPool_DoAddSerializedFile(
+    PyObject* _self, PyObject* serialized_pb) {
+  PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
+  char* buf;
+  Py_ssize_t size;
+  upb_arena* arena = upb_arena_new();
+  if (!arena) PYUPB_RETURN_OOM;
+  PyObject* result = NULL;
 
   if (PyBytes_AsStringAndSize(serialized_pb, &buf, &size) < 0) {
     goto done;
@@ -196,19 +219,45 @@ done:
   return result;
 }
 
-static PyObject* PyUpb_DescriptorPool_Add(PyObject* _self,
-                                          PyObject* file_desc) {
+static PyObject* PyUpb_DescriptorPool_DoAdd(PyObject* _self,
+                                            PyObject* file_desc) {
   PyObject* subargs = PyTuple_New(0);
-  // TODO: check file_desc type more.
+  if (!PyUpb_CMessage_Check(file_desc)) return NULL;
+  const upb_msgdef* m = PyUpb_CMessage_GetMsgdef(file_desc);
+  const char* file_proto_name = "google.protobuf.FileDescriptorProto";
+  if (strcmp(upb_msgdef_fullname(m), file_proto_name) != 0) {
+    fprintf(stderr, "YO: %s\n", upb_msgdef_fullname(m));
+    return PyErr_Format(PyExc_TypeError, "Can only add FileDescriptorProto");
+  }
   PyObject* serialized =
       PyUpb_CMessage_SerializeToString(file_desc, subargs, NULL);
   Py_DECREF(subargs);
   if (!serialized) return NULL;
-  PyObject* ret = PyUpb_DescriptorPool_AddSerializedFile(_self, serialized);
+  PyObject* ret = PyUpb_DescriptorPool_DoAddSerializedFile(_self, serialized);
   Py_DECREF(serialized);
   return ret;
 }
 
+/*
+ * PyUpb_DescriptorPool_AddSerializedFile()
+ *
+ * Implements:
+ *   DescriptorPool.AddSerializedFile(self, serialized_file_descriptor)
+ *
+ * Adds the given serialized FileDescriptorProto to the pool.
+ */
+static PyObject* PyUpb_DescriptorPool_AddSerializedFile(
+    PyObject * self, PyObject * serialized_pb) {
+  if (!PyUpb_DescriptorPool_CheckNoDatabase(self)) return NULL;
+  return PyUpb_DescriptorPool_DoAddSerializedFile(self, serialized_pb);
+}
+
+static PyObject* PyUpb_DescriptorPool_Add(PyObject* self,
+                                          PyObject* file_desc) {
+  if (!PyUpb_DescriptorPool_CheckNoDatabase(self)) return NULL;
+  return PyUpb_DescriptorPool_DoAdd(self, file_desc);
+}
+
 /*
  * PyUpb_DescriptorPool_FindFileByName()
  *
@@ -265,6 +314,10 @@ static PyObject* PyUpb_DescriptorPool_FindMessageTypeByName(PyObject* _self,
   if (!name) return NULL;
 
   const upb_msgdef* m = upb_symtab_lookupmsg(self->symtab, name);
+  if (m == NULL && self->db) {
+    if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
+    m = upb_symtab_lookupmsg(self->symtab, name);
+  }
   if (m == NULL) {
     return PyErr_Format(PyExc_KeyError, "Couldn't find message %.200s", name);
   }

python/pb_unit_tests/descriptor_pool_test_wrapper.py

@@ -37,7 +37,6 @@ descriptor_pool_test.AddDescriptorTest.testService.__unittest_expecting_failure_
 descriptor_pool_test.CreateDescriptorPoolTest.testFindFieldByName.__unittest_expecting_failure__ = True
 descriptor_pool_test.CreateDescriptorPoolTest.testFindService.__unittest_expecting_failure__ = True
 descriptor_pool_test.CreateDescriptorPoolTest.testFindTypeErrors.__unittest_expecting_failure__ = True
-descriptor_pool_test.CreateDescriptorPoolTest.testUserDefinedDB.__unittest_expecting_failure__ = True
 descriptor_pool_test.SecondaryDescriptorFromDescriptorDB.testErrorCollector.__unittest_expecting_failure__ = True
 
 # Some tests are defined in a base class and inherited by multiple sub-classes.
@@ -50,14 +49,12 @@ def wrap(cls, method):
 
 wrap(descriptor_pool_test.CreateDescriptorPoolTest, "testComplexNesting")
 wrap(descriptor_pool_test.DefaultDescriptorPoolTest, "testComplexNesting")
-wrap(descriptor_pool_test.SecondaryDescriptorFromDescriptorDB, "testFindAllExtensions")
 wrap(descriptor_pool_test.SecondaryDescriptorFromDescriptorDB, "testFindEnumTypeByName")
 wrap(descriptor_pool_test.SecondaryDescriptorFromDescriptorDB, "testFindExtensionByName")
 wrap(descriptor_pool_test.SecondaryDescriptorFromDescriptorDB, "testFindExtensionByNumber")
 wrap(descriptor_pool_test.SecondaryDescriptorFromDescriptorDB, "testFindFileByName")
 wrap(descriptor_pool_test.SecondaryDescriptorFromDescriptorDB, "testFindFileContainingSymbol")
 wrap(descriptor_pool_test.SecondaryDescriptorFromDescriptorDB, "testFindOneofByName")
-wrap(descriptor_pool_test.SecondaryDescriptorFromDescriptorDB, "testFindMessageTypeByName")
 
 if __name__ == '__main__':
   unittest.main(module=descriptor_pool_test, verbosity=2)