Change upb python to keep the meta class ref on message descriptor. Also add Garbage Collection on meta class and message descriptor.

python/google/protobuf/internal/message_factory_test.py

@@ -42,7 +42,7 @@
 from google.protobuf import descriptor_database
 from google.protobuf import descriptor_pool
 from google.protobuf import message_factory
-
+from google.protobuf import descriptor
 
 @testing_refleaks.TestCase
 class MessageFactoryTest(unittest.TestCase):
@@ -296,6 +296,25 @@ def testExtensionValueInDifferentFile(self):
     self.assertEqual(234, m.Extensions[ext1].setting)
     self.assertEqual(345, m.Extensions[ext2].setting)
 
+  def testDescriptorKeepConcreteClass(self):
+    def loadFile():
+      f= descriptor_pb2.FileDescriptorProto(
+        name='google/protobuf/internal/meta_class.proto',
+        package='google.protobuf.python.internal')
+      msg_proto = f.message_type.add(name='Empty')
+      msg_proto.nested_type.add(name='Nested')
+      msg_proto.field.add(name='nested_field',
+                          number=1,
+                          label=descriptor.FieldDescriptor.LABEL_REPEATED,
+                          type=descriptor.FieldDescriptor.TYPE_MESSAGE,
+                          type_name='Nested')
+      return message_factory.GetMessages([f])
+
+    messages = loadFile()
+    for des, meta_class in messages.items():
+      message = meta_class()
+      nested_des = message.DESCRIPTOR.nested_types_by_name['Nested']
+      nested_msg = nested_des._concrete_class()
 
 if __name__ == '__main__':
   unittest.main()

upb/python/descriptor.c

@@ -46,9 +46,10 @@
 
 typedef struct {
   PyObject_HEAD;
-  PyObject* pool;     // We own a ref.
-  const void* def;    // Type depends on the class. Kept alive by "pool".
-  PyObject* options;  // NULL if not present or not cached.
+  PyObject* pool;          // We own a ref.
+  const void* def;         // Type depends on the class. Kept alive by "pool".
+  PyObject* options;       // NULL if not present or not cached.
+  PyObject* message_meta;  // We own a ref.
 } PyUpb_DescriptorBase;
 
 PyObject* PyUpb_AnyDescriptor_GetPool(PyObject* desc) {
@@ -71,6 +72,7 @@ static PyUpb_DescriptorBase* PyUpb_DescriptorBase_DoCreate(
   base->pool = PyUpb_DescriptorPool_Get(upb_FileDef_Pool(file));
   base->def = def;
   base->options = NULL;
+  base->message_meta = NULL;
 
   PyUpb_ObjCache_Add(def, &base->ob_base);
   return base;
@@ -194,11 +196,23 @@ static PyObject* PyUpb_DescriptorBase_CopyToProto(PyObject* _self,
 
 static void PyUpb_DescriptorBase_Dealloc(PyUpb_DescriptorBase* base) {
   PyUpb_ObjCache_Delete(base->def);
+  Py_XDECREF(base->message_meta);
   Py_DECREF(base->pool);
   Py_XDECREF(base->options);
   PyUpb_Dealloc(base);
 }
 
+static int PyUpb_Descriptor_Traverse(PyUpb_DescriptorBase* base,
+                                     visitproc visit, void* arg) {
+  Py_VISIT(base->message_meta);
+  return 0;
+}
+
+static int PyUpb_Descriptor_Clear(PyUpb_DescriptorBase* base) {
+  Py_CLEAR(base->message_meta);
+  return 0;
+}
+
 #define DESCRIPTOR_BASE_SLOTS                           \
   {Py_tp_new, (void*)&PyUpb_Forbidden_New}, {           \
     Py_tp_dealloc, (void*)&PyUpb_DescriptorBase_Dealloc \
@@ -219,6 +233,13 @@ PyObject* PyUpb_Descriptor_GetClass(const upb_MessageDef* m) {
   return ret;
 }
 
+void PyUpb_Descriptor_SetClass(PyObject* py_descriptor, PyObject* meta) {
+  PyUpb_DescriptorBase* base = (PyUpb_DescriptorBase*)py_descriptor;
+  Py_XDECREF(base->message_meta);
+  base->message_meta = meta;
+  Py_INCREF(meta);
+}
+
 // The LookupNested*() functions provide name lookup for entities nested inside
 // a message.  This uses the symtab's table, which requires that the symtab is
 // not being mutated concurrently.  We can guarantee this for Python-owned
@@ -631,13 +652,15 @@ static PyType_Slot PyUpb_Descriptor_Slots[] = {
     DESCRIPTOR_BASE_SLOTS,
     {Py_tp_methods, PyUpb_Descriptor_Methods},
     {Py_tp_getset, PyUpb_Descriptor_Getters},
+    {Py_tp_traverse, PyUpb_Descriptor_Traverse},
+    {Py_tp_clear, PyUpb_Descriptor_Clear},
     {0, NULL}};
 
 static PyType_Spec PyUpb_Descriptor_Spec = {
     PYUPB_MODULE_NAME ".Descriptor",  // tp_name
     sizeof(PyUpb_DescriptorBase),     // tp_basicsize
     0,                                // tp_itemsize
-    Py_TPFLAGS_DEFAULT,               // tp_flags
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, // tp_flags
     PyUpb_Descriptor_Slots,
 };
 

upb/python/descriptor.h

@@ -52,6 +52,9 @@ typedef enum {
 // the msgdef |m|, which must be from the same pool.
 PyObject* PyUpb_Descriptor_GetClass(const upb_MessageDef* m);
 
+// Set the message descriptor's meta class.
+void PyUpb_Descriptor_SetClass(PyObject* py_descriptor, PyObject* meta);
+
 // Returns a Python wrapper object for the given def. This will return an
 // existing object if one already exists, otherwise a new object will be
 // created.  The caller always owns a ref on the returned object.

upb/python/message.c

@@ -63,6 +63,8 @@ typedef struct {
   getattrofunc type_getattro;  // PyTypeObject.tp_getattro
   setattrofunc type_setattro;  // PyTypeObject.tp_setattro
   size_t type_basicsize;       // sizeof(PyHeapTypeObject)
+  traverseproc type_traverse;  // PyTypeObject.tp_traverse
+  inquiry type_clear;          // PyTypeObject.tp_clear
 
   // While we can refer to PY_VERSION_HEX in the limited API, this will give us
   // the version of Python we were compiled against, which may be different
@@ -135,6 +137,8 @@ static bool PyUpb_CPythonBits_Init(PyUpb_CPythonBits* bits) {
   bits->type_dealloc = upb_Pre310_PyType_GetDeallocSlot(type);
   bits->type_getattro = PyType_GetSlot(type, Py_tp_getattro);
   bits->type_setattro = PyType_GetSlot(type, Py_tp_setattro);
+  bits->type_traverse = PyType_GetSlot(type, Py_tp_traverse);
+  bits->type_clear = PyType_GetSlot(type, Py_tp_clear);
 
   size = PyObject_GetAttrString((PyObject*)&PyType_Type, "__basicsize__");
   if (!size) goto err;
@@ -145,13 +149,17 @@ static bool PyUpb_CPythonBits_Init(PyUpb_CPythonBits* bits) {
   assert(bits->type_dealloc);
   assert(bits->type_getattro);
   assert(bits->type_setattro);
+  assert(bits->type_traverse);
+  assert(bits->type_clear);
 
 #ifndef Py_LIMITED_API
   assert(bits->type_new == PyType_Type.tp_new);
   assert(bits->type_dealloc == PyType_Type.tp_dealloc);
   assert(bits->type_getattro == PyType_Type.tp_getattro);
   assert(bits->type_setattro == PyType_Type.tp_setattro);
   assert(bits->type_basicsize == sizeof(PyHeapTypeObject));
+  assert(bits->type_traverse == PyType_Type.tp_traverse);
+  assert(bits->type_clear == PyType_Type.tp_clear);
 #endif
 
   sys = PyImport_ImportModule("sys");
@@ -1809,6 +1817,7 @@ PyObject* PyUpb_MessageMeta_DoCreateClass(PyObject* py_descriptor,
   meta->py_message_descriptor = py_descriptor;
   meta->layout = upb_MessageDef_MiniTable(msgdef);
   Py_INCREF(meta->py_message_descriptor);
+  PyUpb_Descriptor_SetClass(py_descriptor, ret);
 
   PyUpb_ObjCache_Add(meta->layout, ret);
 
@@ -1945,10 +1954,23 @@ static PyObject* PyUpb_MessageMeta_GetAttr(PyObject* self, PyObject* name) {
   return NULL;
 }
 
+static int PyUpb_MessageMeta_Traverse(PyObject* self, visitproc visit,
+                                      void* arg) {
+  PyUpb_MessageMeta* meta = PyUpb_GetMessageMeta(self);
+  Py_VISIT(meta->py_message_descriptor);
+  return cpython_bits.type_traverse(self, visit, arg);
+}
+
+static int PyUpb_MessageMeta_Clear(PyObject* self, visitproc visit, void* arg) {
+  return cpython_bits.type_clear(self);
+}
+
 static PyType_Slot PyUpb_MessageMeta_Slots[] = {
     {Py_tp_new, PyUpb_MessageMeta_New},
     {Py_tp_dealloc, PyUpb_MessageMeta_Dealloc},
     {Py_tp_getattro, PyUpb_MessageMeta_GetAttr},
+    {Py_tp_traverse, PyUpb_MessageMeta_Traverse},
+    {Py_tp_clear, PyUpb_MessageMeta_Clear},
     {0, NULL}};
 
 static PyType_Spec PyUpb_MessageMeta_Spec = {
@@ -1957,7 +1979,7 @@ static PyType_Spec PyUpb_MessageMeta_Spec = {
     0,  // tp_itemsize
     // TODO(haberman): remove BASETYPE, Python should just use MessageMeta
     // directly instead of subclassing it.
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,  // tp_flags
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,  // tp_flags
     PyUpb_MessageMeta_Slots,
 };