Merge pull request #473 from haberman/python-map

Added map fields for Python

python/BUILD

@@ -39,6 +39,8 @@ cc_binary(
         "descriptor_containers.h",
         "descriptor_pool.c",
         "descriptor_pool.h",
+        "map.c",
+        "map.h",
         "message.c",
         "message.h",
         "protobuf.c",

python/map.c

@@ -0,0 +1,493 @@
+/*
+ * 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 "python/map.h"
+
+#include "python/convert.h"
+#include "python/message.h"
+#include "python/protobuf.h"
+
+// -----------------------------------------------------------------------------
+// MapContainer
+// -----------------------------------------------------------------------------
+
+typedef struct {
+  PyObject_HEAD
+  PyObject* arena;
+  // The field descriptor (upb_fielddef*).
+  // The low bit indicates whether the container is reified (see ptr below).
+  //   - low bit set: repeated field is a stub (empty map, no underlying data).
+  //   - low bit clear: repeated field is reified (points to upb_array).
+  uintptr_t field;
+  union {
+    PyObject* parent;  // stub: owning pointer to parent message.
+    upb_map* map;      // reified: the data for this array.
+  } ptr;
+  int version;
+} PyUpb_MapContainer;
+
+static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map);
+
+static bool PyUpb_MapContainer_IsStub(PyUpb_MapContainer* self) {
+  return self->field & 1;
+}
+
+// If the map is reified, returns it.  Otherwise, returns NULL.
+// If NULL is returned, the object is empty and has no underlying data.
+static upb_map* PyUpb_MapContainer_GetIfReified(PyUpb_MapContainer* self) {
+  return PyUpb_MapContainer_IsStub(self) ? NULL : self->ptr.map;
+}
+
+static const upb_fielddef* PyUpb_MapContainer_GetField(
+    PyUpb_MapContainer* self) {
+  return (const upb_fielddef*)(self->field & ~(uintptr_t)1);
+}
+
+static void PyUpb_MapContainer_Dealloc(void* _self) {
+  PyUpb_MapContainer* self = _self;
+  Py_DECREF(self->arena);
+  if (PyUpb_MapContainer_IsStub(self)) {
+    PyUpb_CMessage_CacheDelete(self->ptr.parent,
+                               PyUpb_MapContainer_GetField(self));
+    Py_DECREF(self->ptr.parent);
+  } else {
+    PyUpb_ObjCache_Delete(self->ptr.map);
+  }
+  PyUpb_Dealloc(_self);
+}
+
+PyTypeObject* PyUpb_MapContainer_GetClass(const upb_fielddef* f) {
+  assert(upb_fielddef_ismap(f));
+  PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
+  return upb_fielddef_issubmsg(f) ? state->message_map_container_type
+                                  : state->scalar_map_container_type;
+}
+
+PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_fielddef* f,
+                                     PyObject* arena) {
+  PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
+  PyUpb_MapContainer* map = (void*)PyType_GenericAlloc(cls, 0);
+  map->arena = arena;
+  map->field = (uintptr_t)f | 1;
+  map->ptr.parent = parent;
+  map->version = 0;
+  Py_INCREF(arena);
+  Py_INCREF(parent);
+  return &map->ob_base;
+}
+
+void PyUpb_MapContainer_Reify(PyObject* _self, upb_map* map) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  PyUpb_ObjCache_Add(map, &self->ob_base);
+  Py_DECREF(self->ptr.parent);
+  self->ptr.map = map;  // Overwrites self->ptr.parent.
+  self->field &= ~(uintptr_t)1;
+  assert(!PyUpb_MapContainer_IsStub(self));
+}
+
+void PyUpb_MapContainer_Invalidate(PyObject* obj) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)obj;
+  self->version++;
+}
+
+static upb_map* PyUpb_MapContainer_AssureReified(PyUpb_MapContainer* self) {
+  self->version++;
+  upb_map* map = PyUpb_MapContainer_GetIfReified(self);
+  if (map) return map;  // Already writable.
+
+  const upb_fielddef* f = PyUpb_MapContainer_GetField(self);
+  upb_arena* arena = PyUpb_Arena_Get(self->arena);
+  const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
+  const upb_fielddef* key_f = upb_msgdef_field(entry_m, 0);
+  const upb_fielddef* val_f = upb_msgdef_field(entry_m, 1);
+  map = upb_map_new(arena, upb_fielddef_type(key_f), upb_fielddef_type(val_f));
+  upb_msgval msgval = {.map_val = map};
+  PyUpb_CMessage_SetConcreteSubobj(self->ptr.parent, f, msgval);
+  PyUpb_MapContainer_Reify((PyObject*)self, map);
+  return map;
+}
+
+int PyUpb_MapContainer_AssignSubscript(PyObject* _self, PyObject* key,
+                                       PyObject* val) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  upb_map* map = PyUpb_MapContainer_AssureReified(self);
+  const upb_fielddef* f = PyUpb_MapContainer_GetField(self);
+  const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
+  const upb_fielddef* key_f = upb_msgdef_field(entry_m, 0);
+  const upb_fielddef* val_f = upb_msgdef_field(entry_m, 1);
+  upb_arena* arena = PyUpb_Arena_Get(self->arena);
+  upb_msgval u_key, u_val;
+  if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return -1;
+
+  if (val) {
+    if (!PyUpb_PyToUpb(val, val_f, &u_val, arena)) return -1;
+    upb_map_set(map, u_key, u_val, arena);
+  } else {
+    if (!upb_map_delete(map, u_key)) {
+      PyErr_Format(PyExc_KeyError, "Key not present in map");
+      return -1;
+    }
+  }
+  return 0;
+}
+
+PyObject* PyUpb_MapContainer_Subscript(PyObject* _self, PyObject* key) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  upb_map* map = PyUpb_MapContainer_GetIfReified(self);
+  const upb_fielddef* f = PyUpb_MapContainer_GetField(self);
+  const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
+  const upb_fielddef* key_f = upb_msgdef_field(entry_m, 0);
+  const upb_fielddef* val_f = upb_msgdef_field(entry_m, 1);
+  upb_arena* arena = PyUpb_Arena_Get(self->arena);
+  upb_msgval u_key, u_val;
+  if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;
+  if (!map || !upb_map_get(map, u_key, &u_val)) {
+    map = PyUpb_MapContainer_AssureReified(self);
+    upb_arena* arena = PyUpb_Arena_Get(self->arena);
+    if (upb_fielddef_issubmsg(val_f)) {
+      u_val.msg_val = upb_msg_new(upb_fielddef_msgsubdef(val_f), arena);
+    } else {
+      memset(&u_val, 0, sizeof(u_val));
+    }
+    upb_map_set(map, u_key, u_val, arena);
+  }
+  return PyUpb_UpbToPy(u_val, val_f, self->arena);
+}
+
+PyObject* PyUpb_MapContainer_Contains(PyObject* _self, PyObject* key) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  upb_map* map = PyUpb_MapContainer_GetIfReified(self);
+  if (!map) Py_RETURN_FALSE;
+  const upb_fielddef* f = PyUpb_MapContainer_GetField(self);
+  const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
+  const upb_fielddef* key_f = upb_msgdef_field(entry_m, 0);
+  upb_msgval u_key;
+  if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
+  if (upb_map_get(map, u_key, NULL)) {
+    Py_RETURN_TRUE;
+  } else {
+    Py_RETURN_FALSE;
+  }
+}
+
+PyObject* PyUpb_MapContainer_Clear(PyObject* _self, PyObject* key) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  upb_map* map = PyUpb_MapContainer_AssureReified(self);
+  upb_map_clear(map);
+  Py_RETURN_NONE;
+}
+
+static PyObject* PyUpb_MapContainer_Get(PyObject* _self, PyObject* args,
+                                        PyObject* kwargs) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  static const char* kwlist[] = {"key", "default", NULL};
+  PyObject* key;
+  PyObject* default_value = NULL;
+  upb_map* map = PyUpb_MapContainer_GetIfReified(self);
+  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O", (char**)kwlist, &key,
+                                   &default_value)) {
+    return NULL;
+  }
+
+  const upb_fielddef* f = PyUpb_MapContainer_GetField(self);
+  const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
+  const upb_fielddef* key_f = upb_msgdef_field(entry_m, 0);
+  const upb_fielddef* val_f = upb_msgdef_field(entry_m, 1);
+  upb_arena* arena = PyUpb_Arena_Get(self->arena);
+  upb_msgval u_key, u_val;
+  if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;
+  if (map && upb_map_get(map, u_key, &u_val)) {
+    return PyUpb_UpbToPy(u_val, val_f, self->arena);
+  }
+  if (default_value) {
+    Py_INCREF(default_value);
+    return default_value;
+  }
+  Py_RETURN_NONE;
+}
+
+static PyObject* PyUpb_MapContainer_GetEntryClass(PyObject* _self,
+                                                  PyObject* arg) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  const upb_fielddef* f = PyUpb_MapContainer_GetField(self);
+  const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
+  return PyUpb_Descriptor_GetClass(entry_m);
+}
+
+Py_ssize_t PyUpb_MapContainer_Length(PyObject* _self) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  upb_map* map = PyUpb_MapContainer_GetIfReified(self);
+  return map ? upb_map_size(map) : 0;
+}
+
+PyUpb_MapContainer* PyUpb_MapContainer_Check(PyObject* _self) {
+  PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
+  if (!PyObject_TypeCheck(_self, state->message_map_container_type) &&
+      !PyObject_TypeCheck(_self, state->scalar_map_container_type)) {
+    PyErr_Format(PyExc_TypeError, "Expected protobuf map, but got %R", _self);
+    return NULL;
+  }
+  return (PyUpb_MapContainer*)_self;
+}
+
+int PyUpb_CMessage_InitMapAttributes(PyObject* map, PyObject* value,
+                                     const upb_fielddef* f);
+
+static PyObject* PyUpb_MapContainer_MergeFrom(PyObject* _self, PyObject* _arg) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  const upb_fielddef* f = PyUpb_MapContainer_GetField(self);
+
+  if (PyDict_Check(_arg)) {
+    return PyErr_Format(PyExc_AttributeError, "Merging of dict is not allowed");
+  }
+
+  if (PyUpb_CMessage_InitMapAttributes(_self, _arg, f) < 0) {
+    return NULL;
+  }
+
+  Py_RETURN_NONE;
+}
+
+static PyObject* PyUpb_MapContainer_Repr(PyObject* _self) {
+  PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
+  upb_map* map = PyUpb_MapContainer_GetIfReified(self);
+  PyObject* dict = PyDict_New();
+  if (map) {
+    const upb_fielddef* f = PyUpb_MapContainer_GetField(self);
+    const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
+    const upb_fielddef* key_f = upb_msgdef_field(entry_m, 0);
+    const upb_fielddef* val_f = upb_msgdef_field(entry_m, 1);
+    size_t iter = UPB_MAP_BEGIN;
+    while (upb_mapiter_next(map, &iter)) {
+      PyObject* key =
+          PyUpb_UpbToPy(upb_mapiter_key(map, iter), key_f, self->arena);
+      PyObject* val =
+          PyUpb_UpbToPy(upb_mapiter_value(map, iter), val_f, self->arena);
+      if (!key || !val) {
+        Py_XDECREF(key);
+        Py_XDECREF(val);
+        Py_DECREF(dict);
+        return NULL;
+      }
+      PyDict_SetItem(dict, key, val);
+      Py_DECREF(key);
+      Py_DECREF(val);
+    }
+  }
+  PyObject* repr = PyObject_Repr(dict);
+  Py_DECREF(dict);
+  return repr;
+}
+
+PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_map* map,
+                                                const upb_fielddef* f,
+                                                PyObject* arena) {
+  PyUpb_MapContainer* ret = (void*)PyUpb_ObjCache_Get(map);
+  if (ret) return &ret->ob_base;
+
+  PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
+  ret = (void*)PyType_GenericAlloc(cls, 0);
+  ret->arena = arena;
+  ret->field = (uintptr_t)f;
+  ret->ptr.map = map;
+  ret->version = 0;
+  Py_INCREF(arena);
+  PyUpb_ObjCache_Add(map, &ret->ob_base);
+  return &ret->ob_base;
+}
+
+// -----------------------------------------------------------------------------
+// ScalarMapContainer
+// -----------------------------------------------------------------------------
+
+static PyMethodDef PyUpb_ScalarMapContainer_Methods[] = {
+    {"__contains__", PyUpb_MapContainer_Contains, METH_O,
+     "Tests whether a key is a member of the map."},
+    {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
+     "Removes all elements from the map."},
+    {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
+     "Gets the value for the given key if present, or otherwise a default"},
+    {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
+     "Return the class used to build Entries of (key, value) pairs."},
+    {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
+     "Merges a map into the current map."},
+    /*
+   { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
+     "Makes a deep copy of the class." },
+   { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
+     "Outputs picklable representation of the repeated field." },
+   */
+    {NULL, NULL},
+};
+
+static PyType_Slot PyUpb_ScalarMapContainer_Slots[] = {
+    {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
+    {Py_mp_length, PyUpb_MapContainer_Length},
+    {Py_mp_subscript, PyUpb_MapContainer_Subscript},
+    {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
+    {Py_tp_methods, PyUpb_ScalarMapContainer_Methods},
+    {Py_tp_iter, PyUpb_MapIterator_New},
+    {Py_tp_repr, PyUpb_MapContainer_Repr},
+    {Py_tp_hash, PyObject_HashNotImplemented},
+    {0, NULL},
+};
+
+static PyType_Spec PyUpb_ScalarMapContainer_Spec = {
+    PYUPB_MODULE_NAME ".ScalarMapContainer",
+    sizeof(PyUpb_MapContainer),
+    0,
+    Py_TPFLAGS_DEFAULT,
+    PyUpb_ScalarMapContainer_Slots,
+};
+
+// -----------------------------------------------------------------------------
+// MessageMapContainer
+// -----------------------------------------------------------------------------
+
+static PyMethodDef PyUpb_MessageMapContainer_Methods[] = {
+    {"__contains__", PyUpb_MapContainer_Contains, METH_O,
+     "Tests whether the map contains this element."},
+    {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
+     "Removes all elements from the map."},
+    {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
+     "Gets the value for the given key if present, or otherwise a default"},
+    {"get_or_create", PyUpb_MapContainer_Subscript, METH_O,
+     "Alias for getitem, useful to make explicit that the map is mutated."},
+    {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
+     "Return the class used to build Entries of (key, value) pairs."},
+    {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
+     "Merges a map into the current map."},
+    /*
+   { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
+     "Makes a deep copy of the class." },
+   { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
+     "Outputs picklable representation of the repeated field." },
+   */
+    {NULL, NULL},
+};
+
+static PyType_Slot PyUpb_MessageMapContainer_Slots[] = {
+    {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
+    {Py_mp_length, PyUpb_MapContainer_Length},
+    {Py_mp_subscript, PyUpb_MapContainer_Subscript},
+    {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
+    {Py_tp_methods, PyUpb_MessageMapContainer_Methods},
+    {Py_tp_iter, PyUpb_MapIterator_New},
+    {Py_tp_repr, PyUpb_MapContainer_Repr},
+    {Py_tp_hash, PyObject_HashNotImplemented},
+    {0, NULL}};
+
+static PyType_Spec PyUpb_MessageMapContainer_Spec = {
+    PYUPB_MODULE_NAME ".MessageMapContainer", sizeof(PyUpb_MapContainer), 0,
+    Py_TPFLAGS_DEFAULT, PyUpb_MessageMapContainer_Slots};
+
+// -----------------------------------------------------------------------------
+// MapIterator
+// -----------------------------------------------------------------------------
+
+typedef struct {
+  PyObject_HEAD
+  PyUpb_MapContainer* map;  // We own a reference.
+  size_t iter;
+  int version;
+} PyUpb_MapIterator;
+
+static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map) {
+  PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
+  PyUpb_MapIterator* iter =
+      (void*)PyType_GenericAlloc(state->map_iterator_type, 0);
+  iter->map = map;
+  iter->iter = UPB_MAP_BEGIN;
+  iter->version = map->version;
+  Py_INCREF(map);
+  return &iter->ob_base;
+}
+
+static void PyUpb_MapIterator_Dealloc(void* _self) {
+  PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
+  Py_DECREF(&self->map->ob_base);
+  PyUpb_Dealloc(_self);
+}
+
+PyObject* PyUpb_MapIterator_IterNext(PyObject* _self) {
+  PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
+  if (self->version != self->map->version) {
+    return PyErr_Format(PyExc_RuntimeError, "Map modified during iteration.");
+  }
+  upb_map* map = PyUpb_MapContainer_GetIfReified(self->map);
+  if (!map) return NULL;
+  if (!upb_mapiter_next(map, &self->iter)) return NULL;
+  upb_msgval key = upb_mapiter_key(map, self->iter);
+  const upb_fielddef* f = PyUpb_MapContainer_GetField(self->map);
+  const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
+  const upb_fielddef* key_f = upb_msgdef_field(entry_m, 0);
+  return PyUpb_UpbToPy(key, key_f, self->map->arena);
+}
+
+static PyType_Slot PyUpb_MapIterator_Slots[] = {
+    {Py_tp_dealloc, PyUpb_MapIterator_Dealloc},
+    {Py_tp_iter, PyObject_SelfIter},
+    {Py_tp_iternext, PyUpb_MapIterator_IterNext},
+    {0, NULL}};
+
+static PyType_Spec PyUpb_MapIterator_Spec = {
+    PYUPB_MODULE_NAME ".MapIterator", sizeof(PyUpb_MapIterator), 0,
+    Py_TPFLAGS_DEFAULT, PyUpb_MapIterator_Slots};
+
+// -----------------------------------------------------------------------------
+// Top Level
+// -----------------------------------------------------------------------------
+
+static PyObject* GetMutableMappingBase(void) {
+  PyObject* collections = NULL;
+  PyObject* mapping = NULL;
+  PyObject* bases = NULL;
+  if ((collections = PyImport_ImportModule("collections.abc")) &&
+      (mapping = PyObject_GetAttrString(collections, "MutableMapping"))) {
+    bases = Py_BuildValue("(O)", mapping);
+  }
+  Py_XDECREF(collections);
+  Py_XDECREF(mapping);
+  return bases;
+}
+
+bool PyUpb_Map_Init(PyObject* m) {
+  PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
+  PyObject* bases = GetMutableMappingBase();
+  if (!bases) return false;
+
+  state->message_map_container_type =
+      PyUpb_AddClassWithBases(m, &PyUpb_MessageMapContainer_Spec, bases);
+  state->scalar_map_container_type =
+      PyUpb_AddClassWithBases(m, &PyUpb_ScalarMapContainer_Spec, bases);
+  state->map_iterator_type = PyUpb_AddClass(m, &PyUpb_MapIterator_Spec);
+
+  Py_DECREF(bases);
+
+  return state->message_map_container_type &&
+         state->scalar_map_container_type && state->map_iterator_type;
+}

python/map.h

@@ -0,0 +1,61 @@
+/*
+ * 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_MAP_H__
+#define PYUPB_MAP_H__
+
+#include <stdbool.h>
+
+#include "python/python.h"
+#include "upb/def.h"
+
+// Creates a new repeated field stub for field `f` of message object `parent`.
+PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_fielddef* f,
+                                     PyObject* arena);
+
+// Returns a map object wrapping `map`, of field type `f`, which must be on
+// `arena`.  If an existing wrapper object exists, it will be returned,
+// otherwise a new object will be created.  The caller always owns a ref on the
+// returned value.
+PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_map* map,
+                                                const upb_fielddef* f,
+                                                PyObject* arena);
+
+// Reifies a map stub to point to the concrete data in `map`.
+void PyUpb_MapContainer_Reify(PyObject* self, upb_map* map);
+
+// Assigns `self[key] = val` for the map `self`.
+int PyUpb_MapContainer_AssignSubscript(PyObject* self, PyObject* key,
+                                       PyObject* val);
+
+// Invalidates any existing iterators for the map `obj`.
+void PyUpb_MapContainer_Invalidate(PyObject* obj);
+
+// Module-level init.
+bool PyUpb_Map_Init(PyObject* m);
+
+#endif  // PYUPB_MAP_H__

python/message.c

@@ -29,6 +29,7 @@
 
 #include "python/convert.h"
 #include "python/descriptor.h"
+#include "python/map.h"
 #include "python/repeated.h"
 #include "upb/def.h"
 #include "upb/reflection.h"
@@ -203,7 +204,9 @@ bool PyUpb_CMessage_Check(PyObject* self) {
   return true;
 }
 
-upb_msg* PyUpb_CMessage_GetIfWritable(PyObject* _self) {
+// If the message is reified, returns it.  Otherwise, returns NULL.
+// If NULL is returned, the object is empty and has no underlying data.
+upb_msg* PyUpb_CMessage_GetIfReified(PyObject* _self) {
   PyUpb_CMessage* self = (void*)_self;
   return PyUpb_CMessage_IsUnset(self) ? NULL : self->ptr.msg;
 }
@@ -286,7 +289,7 @@ static bool PyUpb_CMessage_InitMessageMapEntry(PyObject* dst, PyObject* src) {
   return true;
 }
 
-static int PyUpb_CMessage_InitMapAttributes(PyObject* map, PyObject* value,
+int PyUpb_CMessage_InitMapAttributes(PyObject* map, PyObject* value,
                                      const upb_fielddef* f) {
   const upb_msgdef* entry_m = upb_fielddef_msgsubdef(f);
   const upb_fielddef* val_f = upb_msgdef_field(entry_m, 1);
@@ -322,18 +325,15 @@ err:
   return ret;
 }
 
-void PyUpb_CMessage_AssureWritable(PyUpb_CMessage* self);
+void PyUpb_CMessage_AssureReified(PyUpb_CMessage* self);
 
 static bool PyUpb_CMessage_InitMapAttribute(PyObject* _self, PyObject* name,
                                             const upb_fielddef* f,
                                             PyObject* value) {
-  // TODO(haberman): disabled until map container is in.
+  PyObject* map = PyUpb_CMessage_GetAttr(_self, name);
-  // PyObject* map = PyUpb_CMessage_GetAttr(_self, name);
+  int ok = PyUpb_CMessage_InitMapAttributes(map, value, f);
-  // int ok = PyUpb_CMessage_InitMapAttributes(map, value, f);
+  Py_DECREF(map);
-  // Py_DECREF(map);
+  return ok >= 0;
-  // return ok >= 0;
-  PyErr_SetString(PyExc_NotImplementedError, "map init");
-  return false;
 }
 
 static bool PyUpb_CMessage_InitRepeatedAttribute(PyObject* _self,
@@ -395,7 +395,7 @@ int PyUpb_CMessage_InitAttributes(PyObject* _self, PyObject* args,
   Py_ssize_t pos = 0;
   PyObject* name;
   PyObject* value;
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
   upb_msg* msg = PyUpb_CMessage_GetMsg(self);
   upb_arena* arena = PyUpb_Arena_Get(self->arena);
 
@@ -437,8 +437,7 @@ static int PyUpb_CMessage_Init(PyObject* _self, PyObject* args,
   return PyUpb_CMessage_InitAttributes(_self, args, kwargs);
 }
 
-static PyObject* PyUpb_CMessage_NewUnset(PyObject* parent,
+static PyObject* PyUpb_CMessage_NewStub(PyObject* parent, const upb_fielddef* f,
-                                         const upb_fielddef* f,
                                         PyObject* arena) {
   const upb_msgdef* sub_m = upb_fielddef_msgsubdef(f);
   PyObject* cls = PyUpb_Descriptor_GetClass(sub_m);
@@ -468,8 +467,8 @@ static bool PyUpb_CMessage_IsEqual(PyUpb_CMessage* m1, PyObject* _m2) {
   const upb_msgdef* m2_msgdef = _PyUpb_CMessage_GetMsgdef(m2);
   assert(m1_msgdef == m2_msgdef);
 #endif
-  const upb_msg* m1_msg = PyUpb_CMessage_GetIfWritable((PyObject*)m1);
+  const upb_msg* m1_msg = PyUpb_CMessage_GetIfReified((PyObject*)m1);
-  const upb_msg* m2_msg = PyUpb_CMessage_GetIfWritable(_m2);
+  const upb_msg* m2_msg = PyUpb_CMessage_GetIfReified(_m2);
   return PyUpb_Message_IsEqual(m1_msg, m2_msg, m1_msgdef);
 }
 
@@ -493,7 +492,7 @@ static void PyUpb_CMessage_SetField(PyUpb_CMessage* parent,
 }
 
 /*
- * PyUpb_CMessage_AssureWritable()
+ * PyUpb_CMessage_AssureReified()
  *
  * This implements the "expando" behavior of Python protos:
  *   foo = FooProto()
@@ -510,7 +509,7 @@ static void PyUpb_CMessage_SetField(PyUpb_CMessage* parent,
  * Post-condition:
  *   PyUpb_CMessage_IsUnset(self) is false
  */
-void PyUpb_CMessage_AssureWritable(PyUpb_CMessage* self) {
+void PyUpb_CMessage_AssureReified(PyUpb_CMessage* self) {
   if (!PyUpb_CMessage_IsUnset(self)) return;
   upb_arena* arena = PyUpb_Arena_Get(self->arena);
 
@@ -541,14 +540,14 @@ void PyUpb_CMessage_AssureWritable(PyUpb_CMessage* self) {
 static void PyUpb_CMessage_SyncSubobjs(PyUpb_CMessage* self);
 
 /*
- * PyUpb_CMessage_SwitchToSet()
+ * PyUpb_CMessage_Reify()
  *
- * The message equivalent of PyUpb_*Container_SwitchToSet(), this transitions
+ * The message equivalent of PyUpb_*Container_Reify(), this transitions
  * the wrapper from the unset state (owning a reference on self->ptr.parent) to the
  * set state (having a non-owning pointer to self->ptr.msg).
  */
-static void PyUpb_CMessage_SwitchToSet(PyUpb_CMessage* self,
+static void PyUpb_CMessage_Reify(PyUpb_CMessage* self, const upb_fielddef* f,
-                                       const upb_fielddef* f, upb_msg* msg) {
+                                 upb_msg* msg) {
   assert(f == PyUpb_CMessage_GetFieldDef(self));
   PyUpb_ObjCache_Add(msg, &self->ob_base);
   Py_DECREF(&self->ptr.parent->ob_base);
@@ -598,15 +597,14 @@ static void PyUpb_CMessage_SyncSubobjs(PyUpb_CMessage* self) {
     PyUpb_WeakMap_DeleteIter(subobj_map, &iter);
     if (upb_fielddef_ismap(f)) {
       if (!msgval.map_val) continue;
-      // TODO(haberman): re-enable when maps are checked in.
+      PyUpb_MapContainer_Reify(obj, (upb_map*)msgval.map_val);
-      // PyUpb_MapContainer_SwitchToSet(obj, (upb_map*)msgval.map_val);
     } else if (upb_fielddef_isseq(f)) {
       if (!msgval.array_val) continue;
       PyUpb_RepeatedContainer_Reify(obj, (upb_array*)msgval.array_val);
     } else {
       PyUpb_CMessage* sub = (void*)obj;
       assert(self == sub->ptr.parent);
-      PyUpb_CMessage_SwitchToSet(sub, f, (upb_msg*)msgval.msg_val);
+      PyUpb_CMessage_Reify(sub, f, (upb_msg*)msgval.msg_val);
     }
   }
 
@@ -659,7 +657,7 @@ void PyUpb_CMessage_CacheDelete(PyObject* _self, const upb_fielddef* f) {
 void PyUpb_CMessage_SetConcreteSubobj(PyObject* _self, const upb_fielddef* f,
                                       upb_msgval subobj) {
   PyUpb_CMessage* self = (void*)_self;
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
   PyUpb_CMessage_CacheDelete(_self, f);
   upb_msg_set(self->ptr.msg, f, subobj, PyUpb_Arena_Get(self->arena));
 }
@@ -730,14 +728,11 @@ PyObject* PyUpb_CMessage_GetUnsetWrapper(PyUpb_CMessage* self,
   if (subobj) return subobj;
 
   if (upb_fielddef_ismap(field)) {
-    // TODO(haberman): re-enable when maps are checked in.
+    subobj = PyUpb_MapContainer_NewStub(_self, field, self->arena);
-    // subobj = PyUpb_MapContainer_NewUnset(_self, field, self->arena);
-    PyErr_SetString(PyExc_NotImplementedError, "unset map");
-    return NULL;
   } else if (upb_fielddef_isseq(field)) {
-    subobj = PyUpb_RepeatedContainer_NewUnset(_self, field, self->arena);
+    subobj = PyUpb_RepeatedContainer_NewStub(_self, field, self->arena);
   } else {
-    subobj = PyUpb_CMessage_NewUnset(&self->ob_base, field, self->arena);
+    subobj = PyUpb_CMessage_NewStub(&self->ob_base, field, self->arena);
   }
   PyUpb_WeakMap_Add(self->unset_subobj_map, field, subobj);
 
@@ -751,12 +746,8 @@ PyObject* PyUpb_CMessage_GetPresentWrapper(PyUpb_CMessage* self,
   upb_mutmsgval mutval =
       upb_msg_mutable(self->ptr.msg, field, PyUpb_Arena_Get(self->arena));
   if (upb_fielddef_ismap(field)) {
-    // TODO(haberman): re-enable when maps are checked in.
+    return PyUpb_MapContainer_GetOrCreateWrapper(mutval.map, field,
-    // return PyUpb_MapContainer_GetOrCreateWrapper(mutval.map, field,
+                                                 self->arena);
-    //                                              self->arena);
-    (void)mutval;
-    PyErr_SetString(PyExc_NotImplementedError, "access map");
-    return NULL;
   } else {
     return PyUpb_RepeatedContainer_GetOrCreateWrapper(mutval.array, field,
                                                       self->arena);
@@ -814,7 +805,7 @@ int PyUpb_CMessage_SetFieldValue(PyObject* _self, const upb_fielddef* field,
     return -1;
   }
 
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
 
   upb_msgval val;
   upb_arena* arena = PyUpb_Arena_Get(self->arena);
@@ -906,7 +897,7 @@ static PyObject* PyUpb_CMessage_HasField(PyObject* _self, PyObject* arg) {
 
 static PyObject* PyUpb_CMessage_ListFields(PyObject* _self, PyObject* arg) {
   PyObject* list = PyList_New(0);
-  upb_msg* msg = PyUpb_CMessage_GetIfWritable(_self);
+  upb_msg* msg = PyUpb_CMessage_GetIfReified(_self);
   if (!msg) return list;
 
   size_t iter1 = UPB_MSG_BEGIN;
@@ -961,7 +952,7 @@ PyObject* PyUpb_CMessage_MergeFrom(PyObject* self, PyObject* arg) {
 
 static PyObject* PyUpb_CMessage_SetInParent(PyObject* _self, PyObject* arg) {
   PyUpb_CMessage* self = (void*)_self;
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
   Py_RETURN_NONE;
 }
 
@@ -988,7 +979,7 @@ PyObject* PyUpb_CMessage_MergeFromString(PyObject* _self, PyObject* arg) {
     return NULL;
   }
 
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
   const upb_msgdef* msgdef = _PyUpb_CMessage_GetMsgdef(self);
   const upb_filedef* file = upb_msgdef_file(msgdef);
   const upb_extreg* extreg = upb_symtab_extreg(upb_filedef_symtab(file));
@@ -1028,7 +1019,7 @@ static PyObject* PyUpb_CMessage_ByteSize(PyObject* self, PyObject* args) {
 }
 
 static PyObject* PyUpb_CMessage_Clear(PyUpb_CMessage* self, PyObject* args) {
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
   const upb_msgdef* msgdef = _PyUpb_CMessage_GetMsgdef(self);
   upb_msg_clear(self->ptr.msg, msgdef);
   Py_RETURN_NONE;
@@ -1036,17 +1027,18 @@ static PyObject* PyUpb_CMessage_Clear(PyUpb_CMessage* self, PyObject* args) {
 
 static void PyUpb_CMessage_AbandonField(PyUpb_CMessage* self,
                                         const upb_fielddef* f) {
-  if (self->unset_subobj_map && upb_fielddef_issubmsg(f)) {
+  if (self->unset_subobj_map && upb_fielddef_issubmsg(f) &&
+      !upb_fielddef_isseq(f)) {
     PyObject* sub = PyUpb_WeakMap_Get(self->unset_subobj_map, f);
     if (sub) {
-      PyUpb_CMessage_AssureWritable((PyUpb_CMessage*)sub);
+      PyUpb_CMessage_AssureReified((PyUpb_CMessage*)sub);
     }
   }
 }
 
 static PyObject* PyUpb_CMessage_ClearExtension(PyUpb_CMessage* self,
                                                PyObject* arg) {
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
   const upb_msgdef* msgdef = _PyUpb_CMessage_GetMsgdef(self);
   const upb_fielddef* f = PyUpb_FieldDescriptor_GetDef(arg);
   if (!f) return NULL;
@@ -1066,7 +1058,7 @@ static PyObject* PyUpb_CMessage_ClearField(PyUpb_CMessage* self,
   //   msg = FooMessage()
   //   msg.foo.Clear()
   //   assert msg.HasField("foo")
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
 
   const upb_fielddef* f;
   const upb_oneofdef* o;
@@ -1089,8 +1081,7 @@ static PyObject* PyUpb_CMessage_ClearField(PyUpb_CMessage* self,
       obj = PyUpb_ObjCache_Get(msgval.map_val);
     }
     if (obj) {
-      // TODO(haberman) add when maps are checked in.
+      PyUpb_MapContainer_Invalidate(obj);
-      // PyUpb_MapContainer_Invalidate(obj);
       Py_DECREF(obj);
     }
   }
@@ -1102,7 +1093,7 @@ static PyObject* PyUpb_CMessage_ClearField(PyUpb_CMessage* self,
 
 static PyObject* PyUpb_CMessage_DiscardUnknownFields(PyUpb_CMessage* self,
                                                      PyObject* arg) {
-  PyUpb_CMessage_AssureWritable(self);
+  PyUpb_CMessage_AssureReified(self);
   const upb_msgdef* msgdef = _PyUpb_CMessage_GetMsgdef(self);
   upb_msg_discardunknown(self->ptr.msg, msgdef, 64);
   Py_RETURN_NONE;
@@ -1111,7 +1102,7 @@ static PyObject* PyUpb_CMessage_DiscardUnknownFields(PyUpb_CMessage* self,
 static PyObject* PyUpb_CMessage_FindInitializationErrors(PyObject* _self,
                                                          PyObject* arg) {
   PyUpb_CMessage* self = (void*)_self;
-  upb_msg* msg = PyUpb_CMessage_GetIfWritable(_self);
+  upb_msg* msg = PyUpb_CMessage_GetIfReified(_self);
   if (!msg) return PyList_New(0);
   const upb_msgdef* msgdef = _PyUpb_CMessage_GetMsgdef(self);
   const upb_symtab* ext_pool = NULL;  // TODO
@@ -1161,7 +1152,7 @@ err:
 
 static PyObject* PyUpb_CMessage_HasExtension(PyObject* _self,
                                              PyObject* ext_desc) {
-  upb_msg* msg = PyUpb_CMessage_GetIfWritable(_self);
+  upb_msg* msg = PyUpb_CMessage_GetIfReified(_self);
   const upb_fielddef* f = PyUpb_FieldDescriptor_GetDef(ext_desc);
   if (!f) return NULL;
   if (!msg) Py_RETURN_FALSE;
@@ -1224,7 +1215,7 @@ static PyObject* PyUpb_CMessage_WhichOneof(PyObject* _self, PyObject* name) {
   if (!PyUpb_CMessage_LookupName(self, name, NULL, &o, PyExc_ValueError)) {
     return NULL;
   }
-  upb_msg* msg = PyUpb_CMessage_GetIfWritable(_self);
+  upb_msg* msg = PyUpb_CMessage_GetIfReified(_self);
   if (!msg) Py_RETURN_NONE;
   const upb_fielddef* f = upb_msg_whichoneof(msg, o);
   if (!f) Py_RETURN_NONE;

python/message.h

@@ -52,9 +52,9 @@ PyObject* PyUpb_CMessage_Get(upb_msg* u_msg, const upb_msgdef* m,
 // returns false on failure.
 bool PyUpb_CMessage_Check(PyObject* self);
 
-// Gets the upb_msg* for this message object if the message is set/writable.
+// Gets the upb_msg* for this message object if the message is reified.
 // Otherwise returns NULL.
-upb_msg* PyUpb_CMessage_GetIfWritable(PyObject* _self);
+upb_msg* PyUpb_CMessage_GetIfReified(PyObject* _self);
 
 // Returns the `upb_msgdef` for a given CMessage.
 const upb_msgdef* PyUpb_CMessage_GetMsgdef(PyObject* self);

python/pb_unit_tests/message_test_wrapper.py

@@ -59,41 +59,14 @@ message_test.Proto2Test.testMergeFromExtensions.__unittest_expecting_failure__ =
 message_test.Proto2Test.testParsingMerge.__unittest_expecting_failure__ = True
 message_test.Proto2Test.testPickleIncompleteProto.__unittest_expecting_failure__ = True
 message_test.Proto2Test.testPythonicInit.__unittest_expecting_failure__ = True
-message_test.Proto2Test.testUnknownEnumMap.__unittest_expecting_failure__ = True
 message_test.Proto2Test.test_documentation.__unittest_expecting_failure__ = True
 message_test.Proto3Test.testCopyFromBadType.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testIntegerMapWithLongs.__unittest_expecting_failure__ = True
 message_test.Proto3Test.testMapAssignmentCausesPresence.__unittest_expecting_failure__ = True
 message_test.Proto3Test.testMapAssignmentCausesPresenceForSubmessages.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapByteSize.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapConstruction.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapDelete.__unittest_expecting_failure__ = True
 message_test.Proto3Test.testMapDeterministicSerialization.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapEntryAlwaysSerialized.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapFieldRaisesCorrectError.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapFindInitializationErrorsSmokeTest.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapGet.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapItems.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapIterInvalidatedByClearField.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapIteration.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapIterationClearMessage.__unittest_expecting_failure__ = True
 message_test.Proto3Test.testMapMergeFrom.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapMessageFieldConstruction.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapScalarFieldConstruction.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapValidAfterFieldCleared.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapsAreMapping.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMapsCompare.__unittest_expecting_failure__ = True
 message_test.Proto3Test.testMergeFrom.__unittest_expecting_failure__ = True
 message_test.Proto3Test.testMergeFromBadType.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMessageMap.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMessageMapItemValidAfterTopMessageCleared.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testMessageMapValidAfterFieldCleared.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testModifyMapWhileIterating.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testNestedMessageMapItemDelete.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testScalarMap.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testScalarMapDefaults.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testStringUnicodeConversionInMap.__unittest_expecting_failure__ = True
-message_test.Proto3Test.testSubmessageMap.__unittest_expecting_failure__ = True
 message_test.OversizeProtosTest.testAssertOversizeProto.__unittest_expecting_failure__ = True
 message_test.OversizeProtosTest.testSucceedOversizeProto.__unittest_expecting_failure__ = True
 

python/pb_unit_tests/unknown_fields_test_wrapper.py

@@ -35,7 +35,6 @@ unknown_fields_test.UnknownFieldsAccessorsTest.testSubUnknownFields.__unittest_e
 unknown_fields_test.UnknownFieldsAccessorsTest.testUnknownExtensions.__unittest_expecting_failure__ = True
 unknown_fields_test.UnknownFieldsAccessorsTest.testUnknownField.__unittest_expecting_failure__ = True
 unknown_fields_test.UnknownFieldsAccessorsTest.testUnknownFieldsNoMemoryLeak.__unittest_expecting_failure__ = True
-unknown_fields_test.UnknownFieldsTest.testDiscardUnknownFields.__unittest_expecting_failure__ = True
 unknown_fields_test.UnknownFieldsTest.testEquals.__unittest_expecting_failure__ = True
 unknown_fields_test.UnknownFieldsTest.testSerializeMessageSetWireFormatUnknownExtension.__unittest_expecting_failure__ = True
 

python/protobuf.c

@@ -30,6 +30,7 @@
 #include "python/descriptor.h"
 #include "python/descriptor_containers.h"
 #include "python/descriptor_pool.h"
+#include "python/map.h"
 #include "python/message.h"
 #include "python/repeated.h"
 
@@ -254,8 +255,22 @@ static const char *PyUpb_GetClassName(PyType_Spec *spec) {
 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
+  if (PyModule_AddObject(m, name, type) < 0) {
-                                                        : NULL;
+    Py_XDECREF(type);
+    return NULL;
+  }
+  return (PyTypeObject*)type;
+}
+
+PyTypeObject* PyUpb_AddClassWithBases(PyObject* m, PyType_Spec* spec,
+                                      PyObject* bases) {
+  PyObject* type = PyType_FromSpecWithBases(spec, bases);
+  const char* name = PyUpb_GetClassName(spec);
+  if (PyModule_AddObject(m, name, type) < 0) {
+    Py_XDECREF(type);
+    return NULL;
+  }
+  return (PyTypeObject*)type;
 }
 
 const char *PyUpb_GetStrData(PyObject *obj) {
@@ -290,7 +305,7 @@ PyMODINIT_FUNC PyInit__message(void) {
   state->obj_cache = PyUpb_WeakMap_New();
 
   if (!PyUpb_InitDescriptorContainers(m) || !PyUpb_InitDescriptorPool(m) ||
-      !PyUpb_InitDescriptor(m) || !PyUpb_InitArena(m) ||
+      !PyUpb_InitDescriptor(m) || !PyUpb_InitArena(m) || !PyUpb_Map_Init(m) ||
       !PyUpb_InitMessage(m) || !PyUpb_Repeated_Init(m)) {
     Py_DECREF(m);
     return NULL;

python/protobuf.h

@@ -64,6 +64,11 @@ typedef struct {
   // From descriptor_pool.c
   PyTypeObject *descriptor_pool_type;
 
+  // From map.c
+  PyTypeObject* map_iterator_type;
+  PyTypeObject* message_map_container_type;
+  PyTypeObject* scalar_map_container_type;
+
   // From message.c
   PyObject *decode_error_class;
   PyObject* descriptor_string;
@@ -165,6 +170,11 @@ 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);
 
+// Like PyUpb_AddClass(), but allows you to specify a tuple of base classes
+// in `bases`.
+PyTypeObject* PyUpb_AddClassWithBases(PyObject* m, PyType_Spec* spec,
+                                      PyObject* bases);
+
 // A function that implements the tp_new slot for types that we do not allow
 // users to create directly. This will immediately fail with an error message.
 PyObject *PyUpb_Forbidden_New(PyObject *cls, PyObject *args, PyObject *kwds);

python/repeated.c

@@ -99,7 +99,9 @@ static const upb_fielddef* PyUpb_RepeatedContainer_GetField(
       PyUpb_RepeatedContainer_GetFieldDescriptor(self));
 }
 
-static upb_array* PyUpb_RepeatedContainer_GetIfWritable(
+// If the repeated field is reified, returns it.  Otherwise, returns NULL.
+// If NULL is returned, the object is empty and has no underlying data.
+static upb_array* PyUpb_RepeatedContainer_GetIfReified(
     PyUpb_RepeatedContainer* self) {
   return PyUpb_RepeatedContainer_IsStub(self) ? NULL : self->ptr.arr;
 }
@@ -114,9 +116,9 @@ void PyUpb_RepeatedContainer_Reify(PyObject* _self, upb_array* arr) {
   assert(!PyUpb_RepeatedContainer_IsStub(self));
 }
 
-static upb_array* PyUpb_RepeatedContainer_AssureWritable(
+static upb_array* PyUpb_RepeatedContainer_AssureReified(
     PyUpb_RepeatedContainer* self) {
-  upb_array* arr = PyUpb_RepeatedContainer_GetIfWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
   if (arr) return arr;  // Already writable.
 
   const upb_fielddef* f = PyUpb_RepeatedContainer_GetField(self);
@@ -151,11 +153,11 @@ static PyTypeObject* PyUpb_RepeatedContainer_GetClass(const upb_fielddef* f) {
 
 static Py_ssize_t PyUpb_RepeatedContainer_Length(PyObject* self) {
   upb_array* arr =
-      PyUpb_RepeatedContainer_GetIfWritable((PyUpb_RepeatedContainer*)self);
+      PyUpb_RepeatedContainer_GetIfReified((PyUpb_RepeatedContainer*)self);
   return arr ? upb_array_size(arr) : 0;
 }
 
-PyObject* PyUpb_RepeatedContainer_NewUnset(PyObject* parent,
+PyObject* PyUpb_RepeatedContainer_NewStub(PyObject* parent,
                                           const upb_fielddef* f,
                                           PyObject* arena) {
   PyTypeObject* cls = PyUpb_RepeatedContainer_GetClass(f);
@@ -187,7 +189,7 @@ PyObject* PyUpb_RepeatedContainer_GetOrCreateWrapper(upb_array* arr,
 
 PyObject* PyUpb_RepeatedContainer_Extend(PyObject* _self, PyObject* value) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
-  upb_array* arr = PyUpb_RepeatedContainer_AssureWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_AssureReified(self);
   size_t start_size = upb_array_size(arr);
   PyObject* it = PyObject_GetIter(value);
   if (!it) {
@@ -223,7 +225,7 @@ PyObject* PyUpb_RepeatedContainer_Extend(PyObject* _self, PyObject* value) {
 static PyObject* PyUpb_RepeatedContainer_Item(PyObject* _self,
                                               Py_ssize_t index) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
-  upb_array* arr = PyUpb_RepeatedContainer_GetIfWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
   Py_ssize_t size = arr ? upb_array_size(arr) : 0;
   if (index < 0 || index >= size) {
     PyErr_Format(PyExc_IndexError, "list index (%zd) out of range", index);
@@ -235,7 +237,7 @@ static PyObject* PyUpb_RepeatedContainer_Item(PyObject* _self,
 
 PyObject* PyUpb_RepeatedContainer_ToList(PyObject* _self) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
-  upb_array* arr = PyUpb_RepeatedContainer_GetIfWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
   if (!arr) return PyList_New(0);
 
   const upb_fielddef* f = PyUpb_RepeatedContainer_GetField(self);
@@ -283,7 +285,7 @@ static PyObject* PyUpb_RepeatedContainer_RichCompare(PyObject* _self,
 static PyObject* PyUpb_RepeatedContainer_Subscript(PyObject* _self,
                                                    PyObject* key) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
-  upb_array* arr = PyUpb_RepeatedContainer_GetIfWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
   Py_ssize_t size = arr ? upb_array_size(arr) : 0;
   Py_ssize_t idx, count, step;
   if (!IndexToRange(key, size, &idx, &count, &step)) return NULL;
@@ -397,7 +399,7 @@ static int PyUpb_RepeatedContainer_AssignSubscript(PyObject* _self,
                                                    PyObject* value) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
   const upb_fielddef* f = PyUpb_RepeatedContainer_GetField(self);
-  upb_array* arr = PyUpb_RepeatedContainer_GetIfWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
   Py_ssize_t size = arr ? upb_array_size(arr) : 0;
   Py_ssize_t idx, count, step;
   if (!IndexToRange(key, size, &idx, &count, &step)) return -1;
@@ -413,7 +415,7 @@ static PyObject* PyUpb_RepeatedContainer_Pop(PyObject* _self, PyObject* args) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
   Py_ssize_t index = -1;
   if (!PyArg_ParseTuple(args, "|n", &index)) return NULL;
-  upb_array* arr = PyUpb_RepeatedContainer_AssureWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_AssureReified(self);
   size_t size = upb_array_size(arr);
   if (index < 0) index += size;
   if (index >= size) index = size - 1;
@@ -427,7 +429,7 @@ static PyObject* PyUpb_RepeatedContainer_Pop(PyObject* _self, PyObject* args) {
 static bool PyUpb_RepeatedContainer_Assign(PyObject* _self, PyObject* list) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
   const upb_fielddef* f = PyUpb_RepeatedContainer_GetField(self);
-  upb_array* arr = PyUpb_RepeatedContainer_AssureWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_AssureReified(self);
   Py_ssize_t size = PyList_Size(list);
   bool submsg = upb_fielddef_issubmsg(f);
   upb_arena* arena = PyUpb_Arena_Get(self->arena);
@@ -435,7 +437,7 @@ static bool PyUpb_RepeatedContainer_Assign(PyObject* _self, PyObject* list) {
     PyObject* obj = PyList_GetItem(list, i);
     upb_msgval msgval;
     if (submsg) {
-      msgval.msg_val = PyUpb_CMessage_GetIfWritable(obj);
+      msgval.msg_val = PyUpb_CMessage_GetIfReified(obj);
       assert(msgval.msg_val);
     } else {
       if (!PyUpb_PyToUpb(obj, f, &msgval, arena)) return false;
@@ -491,7 +493,7 @@ static PyObject* PyUpb_RepeatedContainer_MergeFrom(PyObject* _self,
 
 static PyObject* PyUpb_RepeatedCompositeContainer_AppendNew(PyObject* _self) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
-  upb_array* arr = PyUpb_RepeatedContainer_AssureWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_AssureReified(self);
   if (!arr) return NULL;
   const upb_fielddef* f = PyUpb_RepeatedContainer_GetField(self);
   upb_arena* arena = PyUpb_Arena_Get(self->arena);
@@ -535,7 +537,7 @@ static PyObject* PyUpb_RepeatedContainer_Insert(PyObject* _self,
   Py_ssize_t index;
   PyObject* value;
   if (!PyArg_ParseTuple(args, "nO", &index, &value)) return NULL;
-  upb_array* arr = PyUpb_RepeatedContainer_AssureWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_AssureReified(self);
   if (!arr) return NULL;
 
   // Normalize index.
@@ -620,7 +622,7 @@ static PyType_Spec PyUpb_RepeatedCompositeContainer_Spec = {
 static PyObject* PyUpb_RepeatedScalarContainer_Append(PyObject* _self,
                                                       PyObject* value) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
-  upb_array* arr = PyUpb_RepeatedContainer_AssureWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_AssureReified(self);
   upb_arena* arena = PyUpb_Arena_Get(self->arena);
   const upb_fielddef* f = PyUpb_RepeatedContainer_GetField(self);
   upb_msgval msgval;
@@ -635,7 +637,7 @@ static int PyUpb_RepeatedScalarContainer_AssignItem(PyObject* _self,
                                                     Py_ssize_t index,
                                                     PyObject* item) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
-  upb_array* arr = PyUpb_RepeatedContainer_GetIfWritable(self);
+  upb_array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
   Py_ssize_t size = arr ? upb_array_size(arr) : 0;
   if (index < 0 || index >= size) {
     PyErr_Format(PyExc_IndexError, "list index (%zd) out of range", index);

python/repeated.h

@@ -33,9 +33,8 @@
 #include "python/python.h"
 #include "upb/def.h"
 
-// Creates a new repeated field in the unset state for field `f` of message
+// Creates a new repeated field stub for field `f` of message object `parent`.
-// object `parent`.
+PyObject* PyUpb_RepeatedContainer_NewStub(PyObject* parent,
-PyObject* PyUpb_RepeatedContainer_NewUnset(PyObject* parent,
                                           const upb_fielddef* f,
                                           PyObject* arena);
 
@@ -47,8 +46,7 @@ PyObject* PyUpb_RepeatedContainer_GetOrCreateWrapper(upb_array* arr,
                                                      const upb_fielddef* f,
                                                      PyObject* arena);
 
-// Switches a repeated field in the unset state to be set, with `arr` as the
+// Reifies a repeated field stub to point to the concrete data in `arr`.
-// data being pointed to.
 void PyUpb_RepeatedContainer_Reify(PyObject* self, upb_array* arr);
 
 // Implements repeated_field.extend(iterable).  `_self` must be a repeated

upb/msg.c

@@ -341,6 +341,7 @@ bool _upb_mapsorter_pushmap(_upb_mapsorter *s, upb_descriptortype_t key_type,
       compar = _upb_mapsorter_cmpbool;
       break;
     case UPB_DESCRIPTOR_TYPE_STRING:
+    case UPB_DESCRIPTOR_TYPE_BYTES:
       compar = _upb_mapsorter_cmpstr;
       break;
     default: