Adds implementations for most descriptor objects.

python/BUILD

@@ -45,6 +45,7 @@ cc_binary(
         "descriptor_pool.h",
         "protobuf.c",
         "protobuf.h",
+        "python.h",
     ],
     copts = UPB_DEFAULT_COPTS + [
         # The Python API requires patterns that are ISO C incompatible, like
@@ -64,6 +65,7 @@ cc_binary(
         "//:reflection",
         "//:upb",
         "//upb/util:compare",
+        "//upb/util:def_to_proto",
         "@system_python//:python_headers",
     ],
 )

python/descriptor.h

@@ -30,21 +30,50 @@
 
 #include <stdbool.h>
 
-#include "protobuf.h"
+#include "python/python.h"
 #include "upb/def.h"
 
-PyObject *PyUpb_FieldDescriptor_GetOrCreateWrapper(const upb_fielddef *field,
-                                                   PyObject *pool);
-PyObject *PyUpb_FileDescriptor_GetOrCreateWrapper(const upb_filedef *file,
-                                                  PyObject *pool);
+typedef enum {
+  kPyUpb_Descriptor = 0,
+  kPyUpb_EnumDescriptor = 1,
+  kPyUpb_EnumValueDescriptor = 2,
+  kPyUpb_FieldDescriptor = 3,
+  kPyUpb_FileDescriptor = 4,
+  kPyUpb_MethodDescriptor = 5,
+  kPyUpb_OneofDescriptor = 6,
+  kPyUpb_ServiceDescriptor = 7,
+  kPyUpb_Descriptor_Count = 8,
+} PyUpb_DescriptorType;
 
+// Given a descriptor object |desc|, returns a Python message class object for
+// the msgdef |m|, which must be from the same pool.
+PyObject *PyUpb_Descriptor_GetClass(const upb_msgdef *m);
+
+// 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.
+PyObject *PyUpb_Descriptor_Get(const upb_msgdef *msgdef);
+PyObject *PyUpb_EnumDescriptor_Get(const upb_enumdef *enumdef);
+PyObject *PyUpb_FieldDescriptor_Get(const upb_fielddef *field);
+PyObject *PyUpb_FileDescriptor_Get(const upb_filedef *file);
+PyObject *PyUpb_OneofDescriptor_Get(const upb_oneofdef *oneof);
+PyObject *PyUpb_EnumValueDescriptor_Get(const upb_enumvaldef *enumval);
+PyObject *PyUpb_Descriptor_GetOrCreateWrapper(const upb_msgdef *msg);
+PyObject *PyUpb_ServiceDescriptor_Get(const upb_servicedef *s);
+
+// Returns the underlying |def| for a given wrapper object. The caller must
+// have already verified that the given Python object is of the expected type.
 const upb_filedef *PyUpb_FileDescriptor_GetDef(PyObject *file);
+const upb_fielddef *PyUpb_FieldDescriptor_GetDef(PyObject *file);
+const upb_msgdef *PyUpb_Descriptor_GetDef(PyObject *_self);
+const void *PyUpb_AnyDescriptor_GetDef(PyObject *_self);
 
 // Returns the underlying |def| for a given wrapper object. The caller must
 // have already verified that the given Python object is of the expected type.
 const upb_filedef *PyUpb_FileDescriptor_GetDef(PyObject *file);
 const void *PyUpb_AnyDescriptor_GetDef(PyObject *_self);
 
-bool PyUpb_InitDescriptor(PyObject *m);
+// Module-level init.
+bool PyUpb_InitDescriptor(PyObject* m);
 
-#endif  // PYUPB_DESCRIPTOR_H__
+#endif   // PYUPB_DESCRIPTOR_H__

python/descriptor_pool.c

@@ -41,6 +41,11 @@ typedef struct {
   PyObject* db;
 } PyUpb_DescriptorPool;
 
+PyObject *PyUpb_DescriptorPool_GetDefaultPool() {
+  PyUpb_ModuleState *s = PyUpb_ModuleState_Get();
+  return s->default_pool;
+}
+
 static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type,
                                                PyObject* db) {
   PyUpb_DescriptorPool* pool = PyObject_GC_New(PyUpb_DescriptorPool, type);
@@ -48,9 +53,14 @@ static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type,
   pool->db = db;
   Py_XINCREF(pool->db);
   PyObject_GC_Track(&pool->ob_base);
+  PyUpb_ObjCache_Add(pool->symtab, &pool->ob_base);
   return &pool->ob_base;
 }
 
+upb_symtab *PyUpb_DescriptorPool_GetSymtab(PyObject *pool) {
+  return ((PyUpb_DescriptorPool*)pool)->symtab;
+}
+
 static int PyUpb_DescriptorPool_Traverse(PyUpb_DescriptorPool* self,
                                          visitproc visit, void* arg) {
   Py_VISIT(self->db);
@@ -62,10 +72,17 @@ static int PyUpb_DescriptorPool_Clear(PyUpb_DescriptorPool* self) {
   return 0;
 }
 
+PyObject* PyUpb_DescriptorPool_Get(const upb_symtab *symtab) {
+  PyObject *pool = PyUpb_ObjCache_Get(symtab);
+  assert(pool);
+  return pool;
+}
+
 static void PyUpb_DescriptorPool_Dealloc(PyUpb_DescriptorPool* self) {
   upb_symtab_free(self->symtab);
   PyUpb_DescriptorPool_Clear(self);
-  PyObject_GC_Del(self);
+  PyUpb_ObjCache_Delete(self->symtab);
+  PyUpb_Dealloc(self);
 }
 
 /*
@@ -132,7 +149,7 @@ static PyObject* PyUpb_DescriptorPool_AddSerializedFile(
     goto done;
   }
 
-  result = PyUpb_FileDescriptor_GetOrCreateWrapper(filedef, _self);
+  result = PyUpb_FileDescriptor_Get(filedef);
 
 done:
   upb_arena_free(arena);
@@ -157,7 +174,7 @@ static PyObject* PyUpb_DescriptorPool_FindExtensionByName(PyObject* _self,
     return PyErr_Format(PyExc_KeyError, "Couldn't find extension %.200s", name);
   }
 
-  return PyUpb_FieldDescriptor_GetOrCreateWrapper(field, _self);
+  return PyUpb_FieldDescriptor_Get(field);
 }
 
 static PyMethodDef PyUpb_DescriptorPool_Methods[] = {

python/descriptor_pool.h

@@ -32,6 +32,12 @@
 
 #include "protobuf.h"
 
+PyObject* PyUpb_DescriptorPool_GetSerializedPb(PyObject* _self,
+                                               const char* filename);
+PyObject* PyUpb_DescriptorPool_Get(const upb_symtab *symtab);
+upb_symtab *PyUpb_DescriptorPool_GetSymtab(PyObject *pool);
+PyObject *PyUpb_DescriptorPool_GetDefaultPool(void);
+
 bool PyUpb_InitDescriptorPool(PyObject* m);
 
 #endif  // PYUPB_DESCRIPTOR_POOL_H__

python/protobuf.c

@@ -28,6 +28,7 @@
 #include "protobuf.h"
 
 #include "descriptor.h"
+#include "descriptor_containers.h"
 #include "descriptor_pool.h"
 
 static void PyUpb_ModuleDealloc(void *module) {
@@ -90,6 +91,50 @@ PyObject *PyUpb_ObjCache_Get(const void *key) {
   }
 }
 
+// -----------------------------------------------------------------------------
+// Arena
+// -----------------------------------------------------------------------------
+
+typedef struct {
+  PyObject_HEAD
+  upb_arena *arena;
+} PyUpb_Arena;
+
+PyObject *PyUpb_Arena_New(void) {
+  PyUpb_ModuleState *state = PyUpb_ModuleState_Get();
+  PyUpb_Arena *arena = (void*)PyType_GenericAlloc(state->arena_type, 0);
+  arena->arena = upb_arena_new();
+  return &arena->ob_base;
+}
+
+static void PyUpb_Arena_Dealloc(PyObject *self) {
+  upb_arena_free(PyUpb_Arena_Get(self));
+  PyUpb_Dealloc(self);
+}
+
+upb_arena *PyUpb_Arena_Get(PyObject *arena) {
+  return ((PyUpb_Arena*)arena)->arena;
+}
+
+static PyType_Slot PyUpb_Arena_Slots[] = {
+    {Py_tp_dealloc, PyUpb_Arena_Dealloc},
+    {0, NULL},
+};
+
+static PyType_Spec PyUpb_Arena_Spec = {
+    PYUPB_MODULE_NAME ".Arena",
+    sizeof(PyUpb_Arena),
+    0,  // itemsize
+    Py_TPFLAGS_DEFAULT,
+    PyUpb_Arena_Slots,
+};
+
+static bool PyUpb_InitArena(PyObject *m) {
+  PyUpb_ModuleState *state = PyUpb_ModuleState_GetFromModule(m);
+  state->arena_type = PyUpb_AddClass(m, &PyUpb_Arena_Spec);
+  return state->arena_type;
+}
+
 // -----------------------------------------------------------------------------
 // Utilities
 // -----------------------------------------------------------------------------
@@ -127,6 +172,14 @@ const char *PyUpb_GetStrData(PyObject *obj) {
   }
 }
 
+PyObject *PyUpb_Forbidden_New(PyObject *cls, PyObject *args, PyObject *kwds) {
+  PyObject *name = PyObject_GetAttrString(cls, "__name__");
+  PyErr_Format(PyExc_RuntimeError,
+               "Objects of type %U may not be created directly.", name);
+  Py_XDECREF(name);
+  return NULL;
+}
+
 // -----------------------------------------------------------------------------
 // Module Entry Point
 // -----------------------------------------------------------------------------
@@ -139,7 +192,8 @@ PyMODINIT_FUNC PyInit__message(void) {
   state->obj_cache_arena = upb_arena_new();
   upb_inttable_init(&state->obj_cache, state->obj_cache_arena);
 
-  if (!PyUpb_InitDescriptorPool(m) || !PyUpb_InitDescriptor(m)) {
+  if (!PyUpb_InitDescriptorContainers(m) || !PyUpb_InitDescriptorPool(m) ||
+      !PyUpb_InitDescriptor(m) || !PyUpb_InitArena(m)) {
     Py_DECREF(m);
     return NULL;
   }

python/protobuf.h

@@ -30,15 +30,8 @@
 
 #include <stdbool.h>
 
-#define Py_LIMITED_API 0x03060000
-#include <Python.h>
-
-// This function was not officially added to the limited API until Python 3.10.
-// But in practice it has been stable since Python 3.1.  See:
-//   https://bugs.python.org/issue41784
-PyAPI_FUNC(const char *)
-    PyUnicode_AsUTF8AndSize(PyObject *unicode, Py_ssize_t *size);
-
+#include "python/descriptor.h"
+#include "python/python.h"
 #include "upb/table_internal.h"
 
 #define PYUPB_MODULE_NAME "google.protobuf.pyext._message"
@@ -52,8 +45,7 @@ PyAPI_FUNC(const char *)
 
 typedef struct {
   // From descriptor.c
-  PyTypeObject *field_descriptor_type;
-  PyTypeObject *file_descriptor_type;
+  PyTypeObject *descriptor_types[kPyUpb_Descriptor_Count];
 
   // From descriptor_containers.c
   PyTypeObject *by_name_map_type;
@@ -61,12 +53,16 @@ typedef struct {
   PyTypeObject *descriptor_iterator_type;
   PyTypeObject *generic_sequence_type;
 
+  // From descriptor_pool.c
+  PyObject *default_pool;
+
   // From descriptor_pool.c
   PyTypeObject *descriptor_pool_type;
 
   // From protobuf.c
   upb_arena *obj_cache_arena;
   upb_inttable obj_cache;
+  PyTypeObject *arena_type;
 } PyUpb_ModuleState;
 
 // Returns the global state object from the current interpreter. The current
@@ -95,6 +91,13 @@ void PyUpb_ObjCache_Delete(const void *key);
 // Returns a new reference to an object if it exists, otherwise returns NULL.
 PyObject *PyUpb_ObjCache_Get(const void *key);
 
+// -----------------------------------------------------------------------------
+// Arena
+// -----------------------------------------------------------------------------
+
+PyObject *PyUpb_Arena_New(void);
+upb_arena *PyUpb_Arena_Get(PyObject *arena);
+
 // -----------------------------------------------------------------------------
 // Utilities
 // -----------------------------------------------------------------------------
@@ -104,6 +107,8 @@ PyTypeObject *AddObject(PyObject *m, const char *name, PyType_Spec *spec);
 // Creates a Python type from `spec` and adds it to the given module `m`.
 PyTypeObject *PyUpb_AddClass(PyObject *m, PyType_Spec *spec);
 
+PyObject *PyUpb_Forbidden_New(PyObject *cls, PyObject *args, PyObject *kwds);
+
 // Our standard dealloc func. It follows the guidance defined in:
 //   https://docs.python.org/3/c-api/typeobj.html#c.PyTypeObject.tp_dealloc
 // However it tests Py_TPFLAGS_HEAPTYPE dynamically so that a single dealloc

python/python.h

@@ -0,0 +1,42 @@
+/*
+ * 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 PYUPB_PYTHON_H__
+#define PYUPB_PYTHON_H__
+
+// We restrict ourselves to the limited API, so that we will be ABI-compatible
+// with any version of Python >= 3.6.1  (3.6.1 introduce PySlice_Unpack())
+#define Py_LIMITED_API 0x03060100
+#include <Python.h>
+
+// This function was not officially added to the limited API until Python 3.10.
+// But in practice it has been stable since Python 3.1.  See:
+//   https://bugs.python.org/issue41784
+PyAPI_FUNC(const char *)
+    PyUnicode_AsUTF8AndSize(PyObject *unicode, Py_ssize_t *size);
+
+#endif  // PYUPB_PYTHON_H__