Added descriptor containers.

python/BUILD

@@ -37,6 +37,8 @@ cc_binary(
     srcs = [
         "descriptor.c",
         "descriptor.h",
+        "descriptor_containers.c",
+        "descriptor_containers.h",
         "descriptor_pool.c",
         "descriptor_pool.h",
         "protobuf.c",

python/descriptor.c

@@ -47,6 +47,11 @@ PyObject *PyUpb_AnyDescriptor_GetPool(PyObject *desc) {
   return base->pool;
 }
 
+const void *PyUpb_AnyDescriptor_GetDef(PyObject *desc) {
+  PyUpb_DescriptorBase *base = (void*)desc;
+  return base->def;
+}
+
 static PyObject *PyUpb_DescriptorBase_New(PyTypeObject *subtype, PyObject *args,
                                           PyObject *kwds) {
   return PyErr_Format(PyExc_RuntimeError,

python/descriptor.h

@@ -40,6 +40,11 @@ PyObject *PyUpb_FileDescriptor_GetOrCreateWrapper(const upb_filedef *file,
 
 const upb_filedef *PyUpb_FileDescriptor_GetDef(PyObject *file);
 
+// 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);
 
 #endif  // PYUPB_DESCRIPTOR_H__

python/descriptor_containers.c

@@ -0,0 +1,670 @@
+/*
+ * 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.
+ */
+
+#include "descriptor_containers.h"
+
+#include "upb/def.h"
+
+#include "protobuf.h"
+#include "descriptor.h"
+
+// -----------------------------------------------------------------------------
+// DescriptorIterator
+// -----------------------------------------------------------------------------
+
+typedef struct {
+  PyObject_HEAD
+  const PyUpb_GenericSequence_Funcs *funcs;
+  const void *parent;    // upb_msgdef*, upb_symtab*, etc.
+  PyObject *parent_obj;  // Python object that keeps parent alive, we own a ref.
+  int index;             // Current iterator index.
+} PyUpb_DescriptorIterator;
+
+static void PyUpb_DescriptorIterator_Dealloc(PyUpb_DescriptorIterator *_self) {
+  PyUpb_DescriptorIterator *self = (PyUpb_DescriptorIterator*)_self;
+  Py_DECREF(self->parent_obj);
+  PyUpb_Dealloc(self);
+}
+
+PyObject *PyUpb_DescriptorIterator_New(const PyUpb_GenericSequence_Funcs *funcs,
+                                       const void *parent,
+                                       PyObject *parent_obj) {
+  PyUpb_ModuleState *s = PyUpb_ModuleState_Get();
+  PyUpb_DescriptorIterator *iter =
+      (void *)PyType_GenericAlloc(s->descriptor_iterator_type, 0);
+  iter->funcs = funcs;
+  iter->parent = parent;
+  iter->parent_obj = parent_obj;
+  iter->index = 0;
+  Py_INCREF(iter->parent_obj);
+  return &iter->ob_base;
+}
+
+PyObject* PyUpb_DescriptorIterator_IterNext(PyObject* _self) {
+  PyUpb_DescriptorIterator *self = (PyUpb_DescriptorIterator*)_self;
+  int size = self->funcs->get_elem_count(self->parent);
+  if (self->index >= size) return NULL;
+  const void *elem = self->funcs->index(self->parent, self->index);
+  self->index++;
+  return self->funcs->get_elem_wrapper(elem);
+}
+
+static PyType_Slot PyUpb_DescriptorIterator_Slots[] = {
+  {Py_tp_dealloc, PyUpb_DescriptorIterator_Dealloc},
+  {Py_tp_iter, PyObject_SelfIter},
+  {Py_tp_iternext, PyUpb_DescriptorIterator_IterNext},
+  {0, NULL}
+};
+
+static PyType_Spec PyUpb_DescriptorIterator_Spec = {
+  PYUPB_MODULE_NAME ".DescriptorIterator",   // tp_name
+  sizeof(PyUpb_DescriptorIterator),          // tp_basicsize
+  0,                                         // tp_itemsize
+  Py_TPFLAGS_DEFAULT,                        // tp_flags
+  PyUpb_DescriptorIterator_Slots,
+};
+
+// -----------------------------------------------------------------------------
+// GenericSequence
+// -----------------------------------------------------------------------------
+
+typedef struct {
+  PyObject_HEAD
+  const PyUpb_GenericSequence_Funcs *funcs;
+  const void *parent;    // upb_msgdef*, upb_symtab*, etc.
+  PyObject *parent_obj;  // Python object that keeps parent alive, we own a ref.
+} PyUpb_GenericSequence;
+
+PyObject *PyUpb_GenericSequence_New(
+    const PyUpb_GenericSequence_Funcs *funcs, const void *parent,
+    PyObject *parent_obj) {
+  PyUpb_ModuleState *s = PyUpb_ModuleState_Get();
+  PyUpb_GenericSequence *seq =
+      (PyUpb_GenericSequence *)PyType_GenericAlloc(s->generic_sequence_type, 0);
+  seq->funcs = funcs;
+  seq->parent = parent;
+  seq->parent_obj = parent_obj;
+  Py_INCREF(parent_obj);
+  return &seq->ob_base;
+}
+
+PyUpb_GenericSequence *PyUpb_GenericSequence_Get(PyObject *obj) {
+  assert(Py_TYPE(obj) == PyUpb_ModuleState_Get()->generic_sequence_type);
+  return (PyUpb_GenericSequence*)obj;
+}
+
+static void PyUpb_GenericSequence_Dealloc(PyObject *_self) {
+  PyUpb_GenericSequence *self = PyUpb_GenericSequence_Get(_self);
+  Py_CLEAR(self->parent_obj);
+  PyUpb_Dealloc(self);
+}
+
+static Py_ssize_t PyUpb_GenericSequence_Length(PyObject* _self) {
+  PyUpb_GenericSequence *self = PyUpb_GenericSequence_Get(_self);
+  return self->funcs->get_elem_count(self->parent);
+}
+
+static PyObject *PyUpb_GenericSequence_GetItem(PyObject *_self,
+                                               Py_ssize_t index) {
+  PyUpb_GenericSequence *self = PyUpb_GenericSequence_Get(_self);
+  Py_ssize_t size = self->funcs->get_elem_count(self->parent);
+  if (index < 0 || index >= size) {
+    PyErr_Format(PyExc_IndexError, "list index (%zd) out of range", index);
+    return NULL;
+  }
+  const void *elem = self->funcs->index(self->parent, index);
+  return self->funcs->get_elem_wrapper(elem);
+}
+
+// A sequence container can only be equal to another sequence container, or (for
+// backward compatibility) to a list containing the same items.
+// Returns 1 if equal, 0 if unequal, -1 on error.
+static int PyUpb_GenericSequence_IsEqual(PyUpb_GenericSequence *self,
+                                         PyObject *other) {
+  // Check the identity of C++ pointers.
+  if (PyObject_TypeCheck(other, Py_TYPE(self))) {
+    PyUpb_GenericSequence *other_seq = (void *)other;
+    return self->parent == other_seq->parent && self->funcs == other_seq->funcs;
+  }
+
+  if (!PyList_Check(other)) return 0;
+
+  // return list(self) == other
+  int n = PyUpb_GenericSequence_Length((PyObject*)self);
+  if (n != PyList_Size(other)) {
+    return false;
+  }
+  for (int i = 0; i < n; i++) {
+    PyObject *item1 = PyUpb_GenericSequence_GetItem((PyObject*)self, i);
+    if (!item1) return -1;
+    PyObject* item2 = PyList_GetItem(other, i);
+    if (!item2) return -1;
+    int cmp = PyObject_RichCompareBool(item1, item2, Py_EQ);
+    Py_DECREF(item1);
+    if (cmp != 1) return cmp;
+  }
+  // All items were found and equal
+  return 1;
+}
+
+static PyObject *PyUpb_GenericSequence_RichCompare(PyObject *_self,
+                                                   PyObject *other, int opid) {
+  PyUpb_GenericSequence *self = PyUpb_GenericSequence_Get(_self);
+  if (opid != Py_EQ && opid != Py_NE) {
+    Py_INCREF(Py_NotImplemented);
+    return Py_NotImplemented;
+  }
+  bool ret = (opid == Py_EQ) == PyUpb_GenericSequence_IsEqual(self, other);
+  return PyBool_FromLong(ret);
+}
+
+// Linear search.  Could optimize this in some, cases (defs that have index),
+// not not all (FileDescriptor.dependencies).
+static int PyUpb_GenericSequence_Find(PyObject *_self, PyObject *item) {
+  PyUpb_GenericSequence *self = PyUpb_GenericSequence_Get(_self);
+  const void *item_ptr = PyUpb_AnyDescriptor_GetDef(item);
+  int count = self->funcs->get_elem_count(self->parent);
+  for (int i = 0; i < count; i++) {
+    if (self->funcs->index(self->parent, i) == item_ptr) {
+      return i;
+    }
+  }
+  return -1;
+}
+
+static PyObject* PyUpb_GenericSequence_Index(PyObject* self, PyObject* item) {
+  int position = PyUpb_GenericSequence_Find(self, item);
+  if (position < 0) {
+    PyErr_SetNone(PyExc_ValueError);
+    return NULL;
+  } else {
+    return PyLong_FromLong(position);
+  }
+}
+
+// Implements list.count(): number of occurrences of the item in the sequence.
+// An item can only appear once in a sequence. If it exists, return 1.
+static PyObject *PyUpb_GenericSequence_Count(PyObject *self, PyObject *item) {
+  int position = PyUpb_GenericSequence_Find(self, item);
+  if (position < 0) {
+    return PyLong_FromLong(0);
+  } else {
+    return PyLong_FromLong(1);
+  }
+}
+
+static PyObject *PyUpb_GenericSequence_Append(PyObject *self, PyObject *args) {
+  PyErr_Format(PyExc_TypeError, "'%R' is not a mutable sequence", self);
+  return NULL;
+}
+
+static PyMethodDef PyUpb_GenericSequence_Methods[] = {
+    {"index", PyUpb_GenericSequence_Index, METH_O},
+    {"count", PyUpb_GenericSequence_Count, METH_O},
+    {"append", PyUpb_GenericSequence_Append, METH_O},
+    // This was implemented for Python/C++ but so far has not been required.
+    //{ "__reversed__", (PyCFunction)Reversed, METH_NOARGS, },
+    {NULL}};
+
+static PyType_Slot PyUpb_GenericSequence_Slots[] = {
+  {Py_tp_dealloc, &PyUpb_GenericSequence_Dealloc},
+  {Py_tp_methods, &PyUpb_GenericSequence_Methods},
+  {Py_sq_length, PyUpb_GenericSequence_Length},
+  {Py_sq_item, PyUpb_GenericSequence_GetItem},
+  {Py_tp_richcompare, &PyUpb_GenericSequence_RichCompare},
+  // These were implemented for Python/C++ but so far have not been required.
+  // {Py_tp_repr, &PyUpb_GenericSequence_Repr},
+  // {Py_sq_contains, PyUpb_GenericSequence_Contains},
+  // {Py_mp_subscript, PyUpb_GenericSequence_Subscript},
+  // {Py_mp_ass_subscript, PyUpb_GenericSequence_AssignSubscript},
+  {0, NULL},
+};
+
+static PyType_Spec PyUpb_GenericSequence_Spec = {
+  PYUPB_MODULE_NAME "._GenericSequence", // tp_name
+  sizeof(PyUpb_GenericSequence),         // tp_basicsize
+  0,                                     // tp_itemsize
+  Py_TPFLAGS_DEFAULT,                    // tp_flags
+  PyUpb_GenericSequence_Slots,
+};
+
+// -----------------------------------------------------------------------------
+// ByNameMap
+// -----------------------------------------------------------------------------
+
+typedef struct {
+  PyObject_HEAD
+  const PyUpb_ByNameMap_Funcs *funcs;
+  const void *parent;    // upb_msgdef*, upb_symtab*, etc.
+  PyObject *parent_obj;  // Python object that keeps parent alive, we own a ref.
+} PyUpb_ByNameMap;
+
+PyObject *PyUpb_ByNameMap_New(const PyUpb_ByNameMap_Funcs *funcs,
+                              const void *parent, PyObject *parent_obj) {
+  PyUpb_ModuleState *s = PyUpb_ModuleState_Get();
+  PyUpb_ByNameMap *map = (void*)PyType_GenericAlloc(s->by_name_map_type, 0);
+  map->funcs = funcs;
+  map->parent = parent;
+  map->parent_obj = parent_obj;
+  Py_INCREF(parent_obj);
+  return &map->ob_base;
+}
+
+static void PyUpb_ByNameMap_Dealloc(PyObject *_self) {
+  PyUpb_ByNameMap *self = (void*)_self;
+  Py_DECREF(self->parent_obj);
+  PyUpb_Dealloc(self);
+}
+
+static Py_ssize_t PyUpb_ByNameMap_Length(PyObject* _self) {
+  PyUpb_ByNameMap *self = (void*)_self;
+  return self->funcs->base.get_elem_count(self->parent);
+}
+
+static PyObject *PyUpb_ByNameMap_Subscript(PyObject *_self, PyObject *key) {
+  PyUpb_ByNameMap *self = (void*)_self;
+  const char *name = PyUpb_GetStrData(key);
+  const void *elem = name ? self->funcs->lookup(self->parent, name) : NULL;
+
+  if (elem) {
+    return self->funcs->base.get_elem_wrapper(elem);
+  } else {
+    PyErr_SetObject(PyExc_KeyError, key);
+    return NULL;
+  }
+}
+
+static int PyUpb_ByNameMap_AssignSubscript(PyObject *self, PyObject *key,
+                                           PyObject *value) {
+  PyErr_Format(PyExc_TypeError, PYUPB_MODULE_NAME
+               ".ByNameMap' object does not support item assignment");
+  return -1;
+}
+
+static int PyUpb_ByNameMap_Contains(PyObject *_self, PyObject *key) {
+  PyUpb_ByNameMap *self = (void*)_self;
+  const char *name = PyUpb_GetStrData(key);
+  const void *elem = name ? self->funcs->lookup(self->parent, name) : NULL;
+  return elem ? 1 : 0;
+}
+
+static PyObject *PyUpb_ByNameMap_Get(PyObject *_self, PyObject *args) {
+  PyUpb_ByNameMap *self = (void*)_self;
+  PyObject* key;
+  PyObject* default_value = Py_None;
+  if (!PyArg_UnpackTuple(args, "get", 1, 2, &key, &default_value)) {
+    return NULL;
+  }
+
+  const char *name = PyUpb_GetStrData(key);
+  const void *elem = name ? self->funcs->lookup(self->parent, name) : NULL;
+
+  if (elem) {
+    return self->funcs->base.get_elem_wrapper(elem);
+  } else {
+    Py_INCREF(default_value);
+    return default_value;
+  }
+}
+
+static PyObject *PyUpb_ByNameMap_GetIter(PyObject *_self) {
+  PyUpb_ByNameMap *self = (PyUpb_ByNameMap *)_self;
+  return PyUpb_DescriptorIterator_New(&self->funcs->base, self->parent,
+                                      self->parent_obj);
+}
+
+static PyObject *PyUpb_ByNameMap_Keys(PyObject *_self, PyObject *args) {
+  PyUpb_ByNameMap *self = (PyUpb_ByNameMap *)_self;
+  int n = self->funcs->base.get_elem_count(self->parent);
+  PyObject *ret = PyList_New(n);
+  if (!ret) return NULL;
+  for (int i = 0; i < n; i++) {
+    const void *elem = self->funcs->base.index(self->parent, i);
+    PyObject *key = PyUnicode_FromString(self->funcs->get_elem_name(elem));
+    if (!key) return NULL;
+    PyList_SetItem(ret, i, key);
+  }
+  return ret;
+}
+
+static PyObject *PyUpb_ByNameMap_Values(PyObject *_self, PyObject *args) {
+  PyUpb_ByNameMap *self = (PyUpb_ByNameMap *)_self;
+  int n = self->funcs->base.get_elem_count(self->parent);
+  PyObject *ret = PyList_New(n);
+  if (!ret) return NULL;
+  for (int i = 0; i < n; i++) {
+    const void *elem = self->funcs->base.index(self->parent, i);
+    PyObject *py_elem = self->funcs->base.get_elem_wrapper(elem);
+    if (!elem) return NULL;
+    PyList_SetItem(ret, i, py_elem);
+  }
+  return ret;
+}
+
+static PyObject *PyUpb_ByNameMap_Items(PyObject *_self, PyObject *args) {
+  PyUpb_ByNameMap *self = (PyUpb_ByNameMap *)_self;
+  int n = self->funcs->base.get_elem_count(self->parent);
+  PyObject *ret = PyList_New(n);
+  if (!ret) return NULL;
+  for (int i = 0; i < n; i++) {
+    PyObject *item = PyTuple_New(2);
+    if (!item) return NULL;
+    const void *elem = self->funcs->base.index(self->parent, i);
+    if (!elem) return NULL;
+    PyObject *py_elem = self->funcs->base.get_elem_wrapper(elem);
+    if (!py_elem) return NULL;
+    PyTuple_SetItem(item, 0,
+                    PyUnicode_FromString(self->funcs->get_elem_name(elem)));
+    PyTuple_SetItem(item, 1, py_elem);
+    PyList_SetItem(ret, i, item);
+  }
+  return ret;
+}
+
+// A mapping container can only be equal to another mapping container, or (for
+// backward compatibility) to a dict containing the same items.
+// Returns 1 if equal, 0 if unequal, -1 on error.
+static int PyUpb_ByNameMap_IsEqual(PyUpb_ByNameMap* self, PyObject* other) {
+  // Check the identity of C++ pointers.
+  if (PyObject_TypeCheck(other, Py_TYPE(self))) {
+    PyUpb_ByNameMap *other_map = (void *)other;
+    return self->parent == other_map->parent && self->funcs == other_map->funcs;
+  }
+
+  if (!PyDict_Check(other)) return 0;
+
+  PyObject *self_dict = PyDict_New();
+  PyDict_Merge(self_dict, (PyObject*)self, 0);
+  int eq = PyObject_RichCompareBool(self_dict, other, Py_EQ);
+  Py_DECREF(self_dict);
+  return eq;
+}
+
+static PyObject *PyUpb_ByNameMap_RichCompare(PyObject *_self, PyObject *other,
+                                               int opid) {
+  PyUpb_ByNameMap *self = (void*)_self;
+  if (opid != Py_EQ && opid != Py_NE) {
+    Py_INCREF(Py_NotImplemented);
+    return Py_NotImplemented;
+  }
+  bool ret = (opid == Py_EQ) == PyUpb_ByNameMap_IsEqual(self, other);
+  return PyBool_FromLong(ret);
+}
+
+static PyMethodDef PyUpb_ByNameMap_Methods[] = {
+    {"get", (PyCFunction)&PyUpb_ByNameMap_Get, METH_VARARGS},
+    {"keys", PyUpb_ByNameMap_Keys, METH_NOARGS},
+    {"values", PyUpb_ByNameMap_Values, METH_NOARGS},
+    {"items", PyUpb_ByNameMap_Items, METH_NOARGS},
+    {NULL}};
+
+static PyType_Slot PyUpb_ByNameMap_Slots[] = {
+    {Py_mp_ass_subscript, PyUpb_ByNameMap_AssignSubscript},
+    {Py_mp_length, PyUpb_ByNameMap_Length},
+    {Py_mp_subscript, PyUpb_ByNameMap_Subscript},
+    {Py_sq_contains, &PyUpb_ByNameMap_Contains},
+    {Py_tp_dealloc, &PyUpb_ByNameMap_Dealloc},
+    {Py_tp_iter, PyUpb_ByNameMap_GetIter},
+    {Py_tp_methods, &PyUpb_ByNameMap_Methods},
+    {Py_tp_richcompare, &PyUpb_ByNameMap_RichCompare},
+    {0, NULL},
+};
+
+static PyType_Spec PyUpb_ByNameMap_Spec = {
+  PYUPB_MODULE_NAME "._ByNameMap",       // tp_name
+  sizeof(PyUpb_ByNameMap),               // tp_basicsize
+  0,                                     // tp_itemsize
+  Py_TPFLAGS_DEFAULT,                    // tp_flags
+  PyUpb_ByNameMap_Slots,
+};
+
+// -----------------------------------------------------------------------------
+// ByNumberMap
+// -----------------------------------------------------------------------------
+
+typedef struct {
+  PyObject_HEAD
+  const PyUpb_ByNumberMap_Funcs *funcs;
+  const void *parent;    // upb_msgdef*, upb_symtab*, etc.
+  PyObject *parent_obj;  // Python object that keeps parent alive, we own a ref.
+} PyUpb_ByNumberMap;
+
+PyObject *PyUpb_ByNumberMap_New(const PyUpb_ByNumberMap_Funcs *funcs,
+                              const void *parent, PyObject *parent_obj) {
+  PyUpb_ModuleState *s = PyUpb_ModuleState_Get();
+  PyUpb_ByNumberMap *map = (void*)PyType_GenericAlloc(s->by_number_map_type, 0);
+  map->funcs = funcs;
+  map->parent = parent;
+  map->parent_obj = parent_obj;
+  Py_INCREF(parent_obj);
+  return &map->ob_base;
+}
+
+static void PyUpb_ByNumberMap_Dealloc(PyObject *_self) {
+  PyUpb_ByNumberMap *self = (void*)_self;
+  Py_DECREF(self->parent_obj);
+  PyUpb_Dealloc(self);
+}
+
+static Py_ssize_t PyUpb_ByNumberMap_Length(PyObject* _self) {
+  PyUpb_ByNameMap *self = (void*)_self;
+  return self->funcs->base.get_elem_count(self->parent);
+}
+
+static PyObject *PyUpb_ByNumberMap_Subscript(PyObject *_self, PyObject *key) {
+  PyUpb_ByNumberMap *self = (void*)_self;
+  long num = PyLong_AsLong(key);
+  const void *elem;
+  if (num == -1 && PyErr_Occurred()) {
+    elem = NULL;
+    PyErr_Clear();
+  } else {
+    elem = self->funcs->lookup(self->parent, num);
+  }
+
+  if (elem) {
+    return self->funcs->base.get_elem_wrapper(elem);
+  } else {
+    PyErr_SetObject(PyExc_KeyError, key);
+    return NULL;
+  }
+}
+
+static int PyUpb_ByNumberMap_AssignSubscript(PyObject *self, PyObject *key,
+                                             PyObject *value) {
+  PyErr_Format(PyExc_TypeError, PYUPB_MODULE_NAME
+               ".ByNumberMap' object does not support item assignment");
+  return -1;
+}
+
+static PyObject *PyUpb_ByNumberMap_Get(PyObject *_self, PyObject *args) {
+  PyUpb_ByNumberMap *self = (void*)_self;
+  PyObject* key;
+  PyObject* default_value = Py_None;
+  if (!PyArg_UnpackTuple(args, "get", 1, 2, &key, &default_value)) {
+    return NULL;
+  }
+
+  const void *elem;
+  long num = PyLong_AsLong(key);
+  if (num == -1 && PyErr_Occurred()) {
+    elem = NULL;
+    PyErr_Clear();
+  } else {
+    elem = self->funcs->lookup(self->parent, num);
+  }
+
+  if (elem) {
+    return self->funcs->base.get_elem_wrapper(elem);
+  } else {
+    Py_INCREF(default_value);
+    return default_value;
+  }
+}
+
+static PyObject *PyUpb_ByNumberMap_GetIter(PyObject *_self) {
+  PyUpb_ByNumberMap *self = (PyUpb_ByNumberMap *)_self;
+  return PyUpb_DescriptorIterator_New(&self->funcs->base, self->parent,
+                                      self->parent_obj);
+}
+
+static PyObject *PyUpb_ByNumberMap_Keys(PyObject *_self, PyObject *args) {
+  PyUpb_ByNumberMap *self = (PyUpb_ByNumberMap *)_self;
+  int n = self->funcs->base.get_elem_count(self->parent);
+  PyObject *ret = PyList_New(n);
+  if (!ret) return NULL;
+  for (int i = 0; i < n; i++) {
+    const void *elem = self->funcs->base.index(self->parent, i);
+    PyObject *key = PyLong_FromLong(self->funcs->get_elem_num(elem));
+    if (!key) return NULL;
+    PyList_SetItem(ret, i, key);
+  }
+  return ret;
+}
+
+static PyObject *PyUpb_ByNumberMap_Values(PyObject *_self, PyObject *args) {
+  PyUpb_ByNumberMap *self = (PyUpb_ByNumberMap *)_self;
+  int n = self->funcs->base.get_elem_count(self->parent);
+  PyObject *ret = PyList_New(n);
+  if (!ret) return NULL;
+  for (int i = 0; i < n; i++) {
+    const void *elem = self->funcs->base.index(self->parent, i);
+    if (!elem) return NULL;
+    PyObject *py_elem = self->funcs->base.get_elem_wrapper(elem);
+    if (!py_elem) return NULL;
+    PyList_SetItem(ret, i, py_elem);
+  }
+  return ret;
+}
+
+static PyObject *PyUpb_ByNumberMap_Items(PyObject *_self, PyObject *args) {
+  PyUpb_ByNumberMap *self = (PyUpb_ByNumberMap *)_self;
+  int n = self->funcs->base.get_elem_count(self->parent);
+  PyObject *ret = PyList_New(n);
+  if (!ret) return NULL;
+  for (int i = 0; i < n; i++) {
+    PyObject *item = PyTuple_New(2);
+    if (!item) return NULL;
+    const void *elem = self->funcs->base.index(self->parent, i);
+    if (!elem) return NULL;
+    int number = self->funcs->get_elem_num(elem);
+    PyObject *py_elem = self->funcs->base.get_elem_wrapper(elem);
+    if (!py_elem) return NULL;
+    PyTuple_SetItem(item, 0, PyLong_FromLong(number));
+    PyTuple_SetItem(item, 1, py_elem);
+    PyList_SetItem(ret, i, item);
+  }
+  return ret;
+}
+
+static int PyUpb_ByNumberMap_Contains(PyObject *_self, PyObject *key) {
+  PyUpb_ByNumberMap *self = (PyUpb_ByNumberMap *)_self;
+  long num = PyLong_AsLong(key);
+  const void *elem;
+  if (num == -1 && PyErr_Occurred()) {
+    elem = NULL;
+    PyErr_Clear();
+  } else {
+    elem = self->funcs->lookup(self->parent, num);
+  }
+  return elem ? 1 : 0;
+}
+
+// A mapping container can only be equal to another mapping container, or (for
+// backward compatibility) to a dict containing the same items.
+// Returns 1 if equal, 0 if unequal, -1 on error.
+static int PyUpb_ByNumberMap_IsEqual(PyUpb_ByNumberMap* self, PyObject* other) {
+  // Check the identity of C++ pointers.
+  if (PyObject_TypeCheck(other, Py_TYPE(self))) {
+    PyUpb_ByNumberMap *other_map = (void *)other;
+    return self->parent == other_map->parent && self->funcs == other_map->funcs;
+  }
+
+  if (!PyDict_Check(other)) return 0;
+
+  PyObject *self_dict = PyDict_New();
+  PyDict_Merge(self_dict, (PyObject*)self, 0);
+  int eq = PyObject_RichCompareBool(self_dict, other, Py_EQ);
+  Py_DECREF(self_dict);
+  return eq;
+}
+
+static PyObject *PyUpb_ByNumberMap_RichCompare(PyObject *_self, PyObject *other,
+                                               int opid) {
+  PyUpb_ByNumberMap *self = (void*)_self;
+  if (opid != Py_EQ && opid != Py_NE) {
+    Py_INCREF(Py_NotImplemented);
+    return Py_NotImplemented;
+  }
+  bool ret = (opid == Py_EQ) == PyUpb_ByNumberMap_IsEqual(self, other);
+  return PyBool_FromLong(ret);
+}
+
+static PyMethodDef PyUpb_ByNumberMap_Methods[] = {
+    {"get", (PyCFunction)&PyUpb_ByNumberMap_Get, METH_VARARGS},
+    {"keys", PyUpb_ByNumberMap_Keys, METH_NOARGS},
+    {"values", PyUpb_ByNumberMap_Values, METH_NOARGS},
+    {"items", PyUpb_ByNumberMap_Items, METH_NOARGS},
+    {NULL}};
+
+static PyType_Slot PyUpb_ByNumberMap_Slots[] = {
+    {Py_mp_ass_subscript, PyUpb_ByNumberMap_AssignSubscript},
+    {Py_mp_length, PyUpb_ByNumberMap_Length},
+    {Py_mp_subscript, PyUpb_ByNumberMap_Subscript},
+    {Py_sq_contains, &PyUpb_ByNumberMap_Contains},
+    {Py_tp_dealloc, &PyUpb_ByNumberMap_Dealloc},
+    {Py_tp_iter, PyUpb_ByNumberMap_GetIter},
+    {Py_tp_methods, &PyUpb_ByNumberMap_Methods},
+    {Py_tp_richcompare, &PyUpb_ByNumberMap_RichCompare},
+    {0, NULL},
+};
+
+static PyType_Spec PyUpb_ByNumberMap_Spec = {
+  PYUPB_MODULE_NAME "._ByNumberMap",     // tp_name
+  sizeof(PyUpb_ByNumberMap),             // tp_basicsize
+  0,                                     // tp_itemsize
+  Py_TPFLAGS_DEFAULT,                    // tp_flags
+  PyUpb_ByNumberMap_Slots,
+};
+
+// -----------------------------------------------------------------------------
+// Top Level
+// -----------------------------------------------------------------------------
+
+bool PyUpb_InitDescriptorContainers(PyObject* m) {
+  PyUpb_ModuleState *s = PyUpb_ModuleState_GetFromModule(m);
+
+  s->by_name_map_type = PyUpb_AddClass(m, &PyUpb_ByNameMap_Spec);
+  s->by_number_map_type = PyUpb_AddClass(m, &PyUpb_ByNumberMap_Spec);
+  s->descriptor_iterator_type =
+      PyUpb_AddClass(m, &PyUpb_DescriptorIterator_Spec);
+  s->generic_sequence_type =
+      PyUpb_AddClass(m, &PyUpb_GenericSequence_Spec);
+
+  return s->by_name_map_type && s->by_number_map_type &&
+         s->descriptor_iterator_type && s->generic_sequence_type;
+}

python/descriptor_containers.h

@@ -0,0 +1,103 @@
+/*
+ * 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_DESCRIPTOR_CONTAINERS_H__
+#define PYUPB_DESCRIPTOR_CONTAINERS_H__
+
+// This file defines immutable Python containiner types whose data comes from
+// an underlying descriptor (def).
+//
+// Because there are many instances of these types that vend different kinds of
+// data (fields, oneofs, enums, etc) these types accept a "vtable" of function
+// pointers. This saves us from having to define numerous distinct Python types
+// for each kind of data we want to vend.
+//
+// The underlying upb APIs follow a consistent pattern that allows us to use
+// those functions directly inside these vtables, greatly reducing the amount of
+// "adaptor" code we need to write.
+
+#include <stdbool.h>
+
+#include "upb/def.h"
+
+#include "protobuf.h"
+
+// -----------------------------------------------------------------------------
+// PyUpb_GenericSequence
+// -----------------------------------------------------------------------------
+
+// A Python object that vends a sequence of descriptors.
+
+typedef struct {
+  // Returns the number of elements in the map.
+  int (*get_elem_count)(const void *parent);
+  // Returns an element by index.
+  const void *(*index)(const void *parent, int idx);
+  // Returns a Python object wrapping this element, caller owns a ref.
+  PyObject *(*get_elem_wrapper)(const void *elem);
+} PyUpb_GenericSequence_Funcs;
+
+PyObject *PyUpb_GenericSequence_New(const PyUpb_GenericSequence_Funcs *funcs,
+                                    const void *parent, PyObject *parent_obj);
+
+// -----------------------------------------------------------------------------
+// PyUpb_ByNameMap
+// -----------------------------------------------------------------------------
+
+// A Python object that vends a name->descriptor map.
+
+typedef struct {
+  PyUpb_GenericSequence_Funcs base;
+  // Looks up by name and returns either a pointer to the element or NULL.
+  const void *(*lookup)(const void *parent, const char *key);
+  // Returns the name associated with this element.
+  const char *(*get_elem_name)(const void *elem);
+} PyUpb_ByNameMap_Funcs;
+
+PyObject *PyUpb_ByNameMap_New(const PyUpb_ByNameMap_Funcs *funcs,
+                              const void *parent, PyObject *parent_obj);
+
+// -----------------------------------------------------------------------------
+// PyUpb_ByNumberMap
+// -----------------------------------------------------------------------------
+
+// A Python object that vends a number->descriptor map.
+
+typedef struct {
+  PyUpb_GenericSequence_Funcs base;
+  // Looks up by name and returns either a pointer to the element or NULL.
+  const void *(*lookup)(const void *parent, int num);
+  // Returns the name associated with this element.
+  int (*get_elem_num)(const void *elem);
+} PyUpb_ByNumberMap_Funcs;
+
+PyObject *PyUpb_ByNumberMap_New(const PyUpb_ByNumberMap_Funcs *funcs,
+                                const void *parent, PyObject *parent_obj);
+
+bool PyUpb_InitDescriptorContainers(PyObject* m);
+
+#endif   // PYUPB_DESCRIPTOR_CONTAINERS_H__

python/protobuf.c

@@ -49,9 +49,16 @@ static struct PyModuleDef module_def = {PyModuleDef_HEAD_INIT,
 // ModuleState
 // -----------------------------------------------------------------------------
 
+PyUpb_ModuleState *PyUpb_ModuleState_GetFromModule(PyObject *module) {
+  PyUpb_ModuleState *state = PyModule_GetState(module);
+  assert(state);
+  assert(PyModule_GetDef(module) == &module_def);
+  return state;
+}
+
 PyUpb_ModuleState *PyUpb_ModuleState_Get() {
   PyObject *module = PyState_FindModule(&module_def);
-  return PyModule_GetState(module);
+  return PyUpb_ModuleState_GetFromModule(module);
 }
 
 // -----------------------------------------------------------------------------
@@ -93,6 +100,23 @@ PyTypeObject *AddObject(PyObject *m, const char *name, PyType_Spec *spec) {
                                                         : NULL;
 }
 
+static const char *PyUpb_GetClassName(PyType_Spec *spec) {
+  // spec->name contains a fully-qualified name, like:
+  //   google.protobuf.pyext._message.FooBar
+  //
+  // Find the rightmost '.' to get "FooBar".
+  const char *name = strrchr(spec->name, '.');
+  assert(name);
+  return name + 1;
+}
+
+PyTypeObject *PyUpb_AddClass(PyObject *m, PyType_Spec *spec) {
+  PyObject *type = PyType_FromSpec(spec);
+  const char *name = PyUpb_GetClassName(spec);
+  return type && PyModule_AddObject(m, name, type) == 0 ? (PyTypeObject *)type
+                                                        : NULL;
+}
+
 const char *PyUpb_GetStrData(PyObject *obj) {
   if (PyUnicode_Check(obj)) {
     return PyUnicode_AsUTF8AndSize(obj, NULL);

python/protobuf.h

@@ -55,6 +55,12 @@ typedef struct {
   PyTypeObject *field_descriptor_type;
   PyTypeObject *file_descriptor_type;
 
+  // From descriptor_containers.c
+  PyTypeObject *by_name_map_type;
+  PyTypeObject *by_number_map_type;
+  PyTypeObject *descriptor_iterator_type;
+  PyTypeObject *generic_sequence_type;
+
   // From descriptor_pool.c
   PyTypeObject *descriptor_pool_type;
 
@@ -66,6 +72,7 @@ typedef struct {
 // Returns the global state object from the current interpreter. The current
 // interpreter is looked up from thread-local state.
 PyUpb_ModuleState *PyUpb_ModuleState_Get(void);
+PyUpb_ModuleState *PyUpb_ModuleState_GetFromModule(PyObject *module);
 
 // -----------------------------------------------------------------------------
 // ObjectCache
@@ -93,6 +100,23 @@ PyObject *PyUpb_ObjCache_Get(const void *key);
 // -----------------------------------------------------------------------------
 
 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);
+
+// 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
+// function can work for any type.
+static inline void PyUpb_Dealloc(void *self) {
+  PyTypeObject *tp = Py_TYPE(self);
+  freefunc tp_free = PyType_GetSlot(tp, Py_tp_free);
+  tp_free(self);
+  if (PyType_GetFlags(tp) & Py_TPFLAGS_HEAPTYPE) {
+    Py_DECREF(tp);
+  }
+}
+
 const char *PyUpb_GetStrData(PyObject *obj);
 
 #endif  // PYUPB_PROTOBUF_H__