Merge pull request #461 from haberman/python-message

Message and MessageMeta classes for Python

python/BUILD

@@ -43,6 +43,8 @@ cc_binary(
         "descriptor_containers.h",
         "descriptor_pool.c",
         "descriptor_pool.h",
+        "message.c",
+        "message.h",
         "protobuf.c",
         "protobuf.h",
         "python.h",
@@ -63,9 +65,11 @@ cc_binary(
     deps = [
         ":version_script.lds",
         "//:reflection",
+        "//:textformat",
         "//:upb",
         "//upb/util:compare",
         "//upb/util:def_to_proto",
+        "//upb/util:required_fields",
         "@system_python//:python_headers",
     ],
 )

python/descriptor_pool.c

@@ -46,17 +46,23 @@ PyObject* PyUpb_DescriptorPool_GetDefaultPool() {
   return s->default_pool;
 }
 
-static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type,
-                                               PyObject* db) {
-  PyUpb_DescriptorPool* pool = PyObject_GC_New(PyUpb_DescriptorPool, type);
+static PyObject* PyUpb_DescriptorPool_DoCreateWithCache(PyTypeObject* type,
+                                               PyObject* db,
+                                               PyUpb_WeakMap *obj_cache) {
+  PyUpb_DescriptorPool* pool = (void*)PyType_GenericAlloc(type, 0);
   pool->symtab = upb_symtab_new();
   pool->db = db;
   Py_XINCREF(pool->db);
-  PyObject_GC_Track(&pool->ob_base);
-  PyUpb_ObjCache_Add(pool->symtab, &pool->ob_base);
+  PyUpb_WeakMap_Add(obj_cache, pool->symtab, &pool->ob_base);
   return &pool->ob_base;
 }
 
+static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type,
+                                               PyObject* db) {
+  return PyUpb_DescriptorPool_DoCreateWithCache(type, db,
+                                                PyUpb_ObjCache_Instance());
+}
+
 upb_symtab* PyUpb_DescriptorPool_GetSymtab(PyObject* pool) {
   return ((PyUpb_DescriptorPool*)pool)->symtab;
 }
@@ -235,13 +241,14 @@ static PyType_Spec PyUpb_DescriptorPool_Spec = {
 // -----------------------------------------------------------------------------
 
 bool PyUpb_InitDescriptorPool(PyObject* m) {
+  PyUpb_ModuleState *state = PyUpb_ModuleState_GetFromModule(m);
   PyTypeObject* descriptor_pool_type =
       AddObject(m, "DescriptorPool", &PyUpb_DescriptorPool_Spec);
 
   if (!descriptor_pool_type) return false;
 
-  PyObject* default_pool =
-      PyUpb_DescriptorPool_DoCreate(descriptor_pool_type, NULL);
-  return default_pool &&
-         PyModule_AddObject(m, "default_pool", default_pool) == 0;
+  state->default_pool = PyUpb_DescriptorPool_DoCreateWithCache(
+      descriptor_pool_type, NULL, state->obj_cache);
+  return state->default_pool &&
+         PyModule_AddObject(m, "default_pool", state->default_pool) == 0;
 }

python/message.h

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2009-2021, Google LLC
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of Google LLC nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL Google LLC BE LIABLE FOR ANY DIRECT,
+ * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#ifndef PYPB_MESSAGE_H__
+#define PYPB_MESSAGE_H__
+
+#include <stdbool.h>
+
+#include "python/protobuf.h"
+#include "upb/reflection.h"
+
+// Removes the wrapper object for this field from the unset subobject cache.
+void PyUpb_CMessage_CacheDelete(PyObject* _self, const upb_fielddef* f);
+
+// Sets the field value for `f` to `subobj`, evicting the wrapper object from
+// the "unset subobject" cache now that real data exists for it.  The caller
+// must also update the wrapper associated with `f` to point to `subobj` also.
+void PyUpb_CMessage_SetConcreteSubobj(PyObject* _self, const upb_fielddef* f,
+                                      upb_msgval subobj);
+
+// Gets a Python wrapper object for message `u_msg` of type `m`, returning a
+// cached wrapper if one was previously created.  If a new object is created,
+// it will reference `arena`, which must own `u_msg`.
+PyObject* PyUpb_CMessage_Get(upb_msg* u_msg, const upb_msgdef* m,
+                             PyObject* arena);
+
+// Verifies that a Python object is a message.  Sets a TypeError exception and
+// returns false on failure.
+bool PyUpb_CMessage_Check(PyObject* self);
+
+// Gets the upb_msg* for this message object if the message is set/writable.
+// Otherwise returns NULL.
+upb_msg* PyUpb_CMessage_GetIfWritable(PyObject* _self);
+
+// Returns the `upb_msgdef` for a given CMessage.
+const upb_msgdef* PyUpb_CMessage_GetMsgdef(PyObject* self);
+
+// Functions that match the corresponding methods on the message object.
+PyObject* PyUpb_CMessage_MergeFrom(PyObject* self, PyObject* arg);
+PyObject* PyUpb_CMessage_MergeFromString(PyObject* self, PyObject* arg);
+PyObject* PyUpb_CMessage_SerializeToString(PyObject* self, PyObject* args,
+                                           PyObject* kwargs);
+
+// Sets fields of the message according to the attribuges in `kwargs`.
+int PyUpb_CMessage_InitAttributes(PyObject* _self, PyObject* args,
+                                  PyObject* kwargs);
+
+// Clears the ExtensionDict from the message.  The message must have an
+// ExtensionDict set.
+void PyUpb_CMessage_ClearExtensionDict(PyObject* _self);
+
+// Implements the equivalent of getattr(msg, field), once `field` has
+// already been resolved to a `upb_fielddef*`.
+PyObject* PyUpb_CMessage_GetFieldValue(PyObject* _self,
+                                       const upb_fielddef* field);
+
+// Implements the equivalent of setattr(msg, field, value), once `field` has
+// already been resolved to a `upb_fielddef*`.
+int PyUpb_CMessage_SetFieldValue(PyObject* _self, const upb_fielddef* field,
+                                 PyObject* value);
+
+// Returns the version associated with this message.  The version will be
+// incremented when the message changes.
+int PyUpb_CMessage_GetVersion(PyObject* _self);
+
+// Module-level init.
+bool PyUpb_InitMessage(PyObject* m);
+
+#endif  // PYPB_MESSAGE_H__

python/protobuf.c

@@ -33,7 +33,7 @@
 
 static void PyUpb_ModuleDealloc(void *module) {
   PyUpb_ModuleState *s = PyModule_GetState(module);
-  upb_arena_free(s->obj_cache_arena);
+  PyUpb_WeakMap_Free(s->obj_cache);
 }
 
 static struct PyModuleDef module_def = {PyModuleDef_HEAD_INIT,
@@ -68,40 +68,121 @@ PyUpb_ModuleState *PyUpb_ModuleState_Get(void) {
   return PyUpb_ModuleState_GetFromModule(module);
 }
 
+PyObject *PyUpb_GetWktBases(PyUpb_ModuleState *state) {
+  if (!state->wkt_bases) {
+    PyObject *wkt_module =
+        PyImport_ImportModule("google.protobuf.internal.well_known_types");
+
+    if (wkt_module == NULL) {
+      return false;
+    }
+
+    state->wkt_bases = PyObject_GetAttrString(wkt_module, "WKTBASES");
+    PyObject *m = PyState_FindModule(&module_def);
+    // Reparent ownership to m.
+    PyModule_AddObject(m, "__internal_wktbases", state->wkt_bases);
+    Py_DECREF(wkt_module);
+  }
+
+  return state->wkt_bases;
+}
+
 // -----------------------------------------------------------------------------
-// ObjectCache
+// WeakMap
 // -----------------------------------------------------------------------------
 
+struct PyUpb_WeakMap {
+  upb_inttable table;
+  upb_arena *arena;
+};
+
+PyUpb_WeakMap *PyUpb_WeakMap_New(void) {
+  upb_arena *arena = upb_arena_new();
+  PyUpb_WeakMap *map = upb_arena_malloc(arena, sizeof(*map));
+  map->arena = arena;
+  upb_inttable_init(&map->table, map->arena);
+  return map;
+}
+
+void PyUpb_WeakMap_Free(PyUpb_WeakMap *map) {
+  upb_arena_free(map->arena);
+}
+
+uintptr_t PyUpb_WeakMap_GetKey(const void *key) {
+  uintptr_t n = (uintptr_t)key;
+  assert((n & 7) == 0);
+  return n >> 3;
+}
+
+void PyUpb_WeakMap_Add(PyUpb_WeakMap *map, const void *key, PyObject *py_obj) {
+  upb_inttable_insert(&map->table, PyUpb_WeakMap_GetKey(key),
+                      upb_value_ptr(py_obj), map->arena);
+}
+
+void PyUpb_WeakMap_Delete(PyUpb_WeakMap *map, const void *key) {
+  upb_value val;
+  bool removed =
+      upb_inttable_remove(&map->table, PyUpb_WeakMap_GetKey(key), &val);
+  (void)removed;
+  assert(removed);
+}
+
+void PyUpb_WeakMap_TryDelete(PyUpb_WeakMap *map, const void *key) {
+  upb_inttable_remove(&map->table, PyUpb_WeakMap_GetKey(key), NULL);
+}
+
+PyObject *PyUpb_WeakMap_Get(PyUpb_WeakMap *map, const void *key) {
+  upb_value val;
+  if (upb_inttable_lookup(&map->table, PyUpb_WeakMap_GetKey(key), &val)) {
+    PyObject *ret = upb_value_getptr(val);
+    Py_INCREF(ret);
+    return ret;
+  } else {
+    return NULL;
+  }
+}
+
+bool PyUpb_WeakMap_Next(PyUpb_WeakMap *map, const void **key, PyObject **obj,
+                        intptr_t *iter) {
+  uintptr_t u_key;
+  upb_value val;
+  if (!upb_inttable_next2(&map->table, &u_key, &val, iter)) return false;
+  *key = (void *)(u_key << 3);
+  *obj = upb_value_getptr(val);
+  return true;
+}
+
+void PyUpb_WeakMap_DeleteIter(PyUpb_WeakMap *map, intptr_t *iter) {
+  upb_inttable_removeiter(&map->table, iter);
+}
+
+// -----------------------------------------------------------------------------
+// ObjCache
+// -----------------------------------------------------------------------------
+
+PyUpb_WeakMap *PyUpb_ObjCache_Instance(void) {
+  PyUpb_ModuleState *state = PyUpb_ModuleState_Get();
+  return state->obj_cache;
+}
+
 void PyUpb_ObjCache_Add(const void *key, PyObject *py_obj) {
-  PyUpb_ModuleState *s = PyUpb_ModuleState_Get();
-  upb_inttable_insert(&s->obj_cache, (uintptr_t)key, upb_value_ptr(py_obj),
-                      s->obj_cache_arena);
+  PyUpb_WeakMap_Add(PyUpb_ObjCache_Instance(), key, py_obj);
 }
 
 void PyUpb_ObjCache_Delete(const void *key) {
-  PyUpb_ModuleState *s = PyUpb_ModuleState_MaybeGet();
-  if (!s) {
+  PyUpb_ModuleState *state = PyUpb_ModuleState_MaybeGet();
+  if (!state) {
     // During the shutdown sequence, our object's Dealloc() methods can be
     // called *after* our module Dealloc() method has been called.  At that
     // point our state will be NULL and there is nothing to delete out of the
     // map.
     return;
   }
-  upb_value val;
-  upb_inttable_remove(&s->obj_cache, (uintptr_t)key, &val);
-  assert(upb_value_getptr(val));
+  PyUpb_WeakMap_Delete(state->obj_cache, key);
 }
 
 PyObject *PyUpb_ObjCache_Get(const void *key) {
-  PyUpb_ModuleState *s = PyUpb_ModuleState_Get();
-  upb_value val;
-  if (upb_inttable_lookup(&s->obj_cache, (uintptr_t)key, &val)) {
-    PyObject *ret = upb_value_getptr(val);
-    Py_INCREF(ret);
-    return ret;
-  } else {
-    return NULL;
-  }
+  return PyUpb_WeakMap_Get(PyUpb_ObjCache_Instance(), key);
 }
 
 // -----------------------------------------------------------------------------
@@ -199,11 +280,12 @@ PyObject *PyUpb_Forbidden_New(PyObject *cls, PyObject *args, PyObject *kwds) {
 
 PyMODINIT_FUNC PyInit__message(void) {
   PyObject *m = PyModule_Create(&module_def);
-  PyState_AddModule(m, &module_def);
-  PyUpb_ModuleState *state = PyUpb_ModuleState_Get();
+  if (!m) return NULL;
+
+  PyUpb_ModuleState *state = PyUpb_ModuleState_GetFromModule(m);
 
-  state->obj_cache_arena = upb_arena_new();
-  upb_inttable_init(&state->obj_cache, state->obj_cache_arena);
+  state->wkt_bases = NULL;
+  state->obj_cache = PyUpb_WeakMap_New();
 
   if (!PyUpb_InitDescriptorContainers(m) || !PyUpb_InitDescriptorPool(m) ||
       !PyUpb_InitDescriptor(m) || !PyUpb_InitArena(m)) {

python/protobuf.h

@@ -38,6 +38,9 @@
 
 #define PYUPB_RETURN_OOM return PyErr_SetNone(PyExc_MemoryError), NULL
 
+struct PyUpb_WeakMap;
+typedef struct PyUpb_WeakMap PyUpb_WeakMap;
+
 // -----------------------------------------------------------------------------
 // ModuleState
 // -----------------------------------------------------------------------------
@@ -61,10 +64,19 @@ typedef struct {
   // From descriptor_pool.c
   PyTypeObject *descriptor_pool_type;
 
+  // From message.c
+  PyObject *decode_error_class;
+  PyObject* descriptor_string;
+  PyObject *encode_error_class;
+  PyObject *enum_type_wrapper_class;
+  PyObject *message_class;
+  PyTypeObject *cmessage_type;
+  PyTypeObject *message_meta_type;
+
   // From protobuf.c
-  upb_arena *obj_cache_arena;
-  upb_inttable obj_cache;
+  PyObject *wkt_bases;
   PyTypeObject *arena_type;
+  PyUpb_WeakMap *obj_cache;
 } PyUpb_ModuleState;
 
 // Returns the global state object from the current interpreter. The current
@@ -72,26 +84,66 @@ typedef struct {
 PyUpb_ModuleState *PyUpb_ModuleState_Get(void);
 PyUpb_ModuleState *PyUpb_ModuleState_GetFromModule(PyObject *module);
 
+// Returns NULL if module state is not yet available (during startup).
+// Any use of the module state during startup needs to be passed explicitly.
+PyUpb_ModuleState* PyUpb_ModuleState_MaybeGet(void);
+
+// Returns:
+//   from google.protobuf.internal.well_known_types import WKTBASES
+//
+// This has to be imported lazily rather than at module load time, because
+// otherwise it would cause a circular import.
+PyObject *PyUpb_GetWktBases(PyUpb_ModuleState *state);
+
 // -----------------------------------------------------------------------------
-// ObjectCache
+// WeakMap
 // -----------------------------------------------------------------------------
 
-// The ObjectCache is a weak map that maps C pointers to the corresponding
-// Python wrapper object. We want a consistent Python wrapper object for each
-// C object, both to save memory and to provide object stability (ie. x is x).
+// A WeakMap maps C pointers to the corresponding Python wrapper object. We
+// want a consistent Python wrapper object for each C object, both to save
+// memory and to provide object stability (ie. x is x).
 //
 // Each wrapped object should add itself to the map when it is constructed and
 // remove itself from the map when it is destroyed. The map is weak so it does
 // not take references to the cached objects.
 
-// Adds the given object to the cache, indexed by the given key.
-void PyUpb_ObjCache_Add(const void *key, PyObject *py_obj);
+PyUpb_WeakMap *PyUpb_WeakMap_New(void);
+void PyUpb_WeakMap_Free(PyUpb_WeakMap *map);
+
+// Adds the given object to the map, indexed by the given key.
+void PyUpb_WeakMap_Add(PyUpb_WeakMap *map, const void *key, PyObject *py_obj);
 
 // Removes the given key from the cache. It must exist in the cache currently.
-void PyUpb_ObjCache_Delete(const void *key);
+void PyUpb_WeakMap_Delete(PyUpb_WeakMap *map, const void *key);
+void PyUpb_WeakMap_TryDelete(PyUpb_WeakMap *map, const void *key);
 
 // Returns a new reference to an object if it exists, otherwise returns NULL.
-PyObject *PyUpb_ObjCache_Get(const void *key);
+PyObject *PyUpb_WeakMap_Get(PyUpb_WeakMap *map, const void *key);
+
+#define PYUPB_WEAKMAP_BEGIN UPB_INTTABLE_BEGIN
+
+// Iteration over the weak map, eg.
+//
+// intptr_t it = PYUPB_WEAKMAP_BEGIN;
+// while (PyUpb_WeakMap_Next(map, &key, &obj, &it)) {
+//   // ...
+// }
+//
+// Note that the callee does not own a ref on the returned `obj`.
+bool PyUpb_WeakMap_Next(PyUpb_WeakMap *map, const void **key, PyObject **obj,
+                        intptr_t *iter);
+void PyUpb_WeakMap_DeleteIter(PyUpb_WeakMap *map, intptr_t *iter);
+
+// -----------------------------------------------------------------------------
+// ObjCache
+// -----------------------------------------------------------------------------
+
+// The object cache is a global WeakMap for mapping upb objects to the
+// corresponding wrapper.
+void PyUpb_ObjCache_Add(const void *key, PyObject *py_obj);
+void PyUpb_ObjCache_Delete(const void *key);
+PyObject *PyUpb_ObjCache_Get(const void *key);  // returns NULL if not present.
+PyUpb_WeakMap *PyUpb_ObjCache_Instance(void);
 
 // -----------------------------------------------------------------------------
 // Arena