diff --git a/python/google/protobuf/internal/message_factory_test.py b/python/google/protobuf/internal/message_factory_test.py index de233f03d268..ea3141dd1723 100644 --- a/python/google/protobuf/internal/message_factory_test.py +++ b/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() diff --git a/upb/python/descriptor.c b/upb/python/descriptor.c index 391e864bc5be..9cce8c77512d 100644 --- a/upb/python/descriptor.c +++ b/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, }; diff --git a/upb/python/descriptor.h b/upb/python/descriptor.h index 63a94f8a7e6d..7fa0164f0b32 100644 --- a/upb/python/descriptor.h +++ b/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. diff --git a/upb/python/message.c b/upb/python/message.c index c35e48fdf476..7c8ddebce1d4 100644 --- a/upb/python/message.c +++ b/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,6 +149,8 @@ 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); @@ -152,6 +158,8 @@ static bool PyUpb_CPythonBits_Init(PyUpb_CPythonBits* bits) { 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, };