On demand create message meta class for upb python

PiperOrigin-RevId: 565733360

python/google/protobuf/internal/factory_test1.proto

@@ -29,6 +29,7 @@ message Factory1Message {
   optional NestedFactory1Message nested_factory_1_message = 3;
   optional int32 scalar_value = 4;
   repeated string list_value = 5;
+  map<string, string> map_field = 6;
 
   extensions 1000 to max;
 }

python/google/protobuf/internal/message_factory_test.py

@@ -10,6 +10,7 @@
 __author__ = 'matthewtoia@google.com (Matt Toia)'
 
 import unittest
+import gc
 
 from google.protobuf import descriptor_pb2
 from google.protobuf.internal import api_implementation
@@ -293,5 +294,29 @@ def loadFile():
       nested_des = message.DESCRIPTOR.nested_types_by_name['Nested']
       nested_msg = nested_des._concrete_class()
 
+  def testOndemandCreateMetaClass(self):
+    def loadFile():
+      f = descriptor_pb2.FileDescriptorProto.FromString(
+        factory_test1_pb2.DESCRIPTOR.serialized_pb)
+      return message_factory.GetMessages([f])
+
+    messages = loadFile()
+    data = factory_test1_pb2.Factory1Message()
+    data.map_field['hello'] = 'welcome'
+    # Force GC to collect. UPB python will clean up the map entry class.
+    # cpp extension and pure python will still keep the map entry class.
+    gc.collect()
+    message = messages['google.protobuf.python.internal.Factory1Message']()
+    message.ParseFromString(data.SerializeToString())
+    value = message.map_field
+    values = [
+        # The entry class will be created on demand in upb python.
+        value.GetEntryClass()(key=k, value=value[k]) for k in sorted(value)
+    ]
+    gc.collect()
+    self.assertEqual(1, len(values))
+    self.assertEqual('hello', values[0].key)
+    self.assertEqual('welcome', values[0].value)
+
 if __name__ == '__main__':
   unittest.main()

upb/python/descriptor.c

@@ -232,6 +232,25 @@ PyObject* PyUpb_Descriptor_Get(const upb_MessageDef* m) {
 
 PyObject* PyUpb_Descriptor_GetClass(const upb_MessageDef* m) {
   PyObject* ret = PyUpb_ObjCache_Get(upb_MessageDef_MiniTable(m));
+  if (ret) return ret;
+
+  // On demand create the clss if not exist. However, if users repeatedly
+  // create and destroy a class, it could trigger a loop. This is not an
+  // issue now, but if we see CPU waste for repeatedly create and destroy
+  // in the future, we could make PyUpb_Descriptor_Get() append the descriptor
+  // to an internal list in DescriptorPool, let the pool keep descriptors alive.
+  PyObject* py_descriptor = PyUpb_Descriptor_Get(m);
+  if (py_descriptor == NULL) return NULL;
+  const char* name = upb_MessageDef_Name(m);
+  PyObject* dict = PyDict_New();
+  if (dict == NULL) goto err;
+  int status = PyDict_SetItemString(dict, "DESCRIPTOR", py_descriptor);
+  if (status < 0) goto err;
+  ret = PyUpb_MessageMeta_DoCreateClass(py_descriptor, name, dict);
+
+err:
+  Py_XDECREF(py_descriptor);
+  Py_XDECREF(dict);
   return ret;
 }
 
@@ -503,7 +522,7 @@ static PyObject* PyUpb_Descriptor_GetFullName(PyObject* self, void* closure) {
 static PyObject* PyUpb_Descriptor_GetConcreteClass(PyObject* self,
                                                    void* closure) {
   const upb_MessageDef* msgdef = PyUpb_Descriptor_GetDef(self);
-  return PyUpb_Descriptor_GetClass(msgdef);
+  return PyUpb_ObjCache_Get(upb_MessageDef_MiniTable(msgdef));
 }
 
 static PyObject* PyUpb_Descriptor_GetFile(PyObject* self, void* closure) {

upb/python/message.h

@@ -96,6 +96,10 @@ PyObject* PyUpb_Message_GetFieldValue(PyObject* _self,
 int PyUpb_Message_SetFieldValue(PyObject* _self, const upb_FieldDef* field,
                                 PyObject* value, PyObject* exc);
 
+// Creates message meta class.
+PyObject* PyUpb_MessageMeta_DoCreateClass(PyObject* py_descriptor,
+                                          const char* name, PyObject* dict);
+
 // Returns the version associated with this message.  The version will be
 // incremented when the message changes.
 int PyUpb_Message_GetVersion(PyObject* _self);