blob: eea57e63d45364ea6caf1078ae27660cbd1d1cf6 [file] [log] [blame]
Joshua Habermanbf74b3e2021-12-30 00:52:12 -08001/*
2 * Copyright (c) 2009-2021, Google LLC
3 * All rights reserved.
4 *
5 * Redistribution and use in source and binary forms, with or without
6 * modification, are permitted provided that the following conditions are met:
7 * * Redistributions of source code must retain the above copyright
8 * notice, this list of conditions and the following disclaimer.
9 * * Redistributions in binary form must reproduce the above copyright
10 * notice, this list of conditions and the following disclaimer in the
11 * documentation and/or other materials provided with the distribution.
12 * * Neither the name of Google LLC nor the
13 * names of its contributors may be used to endorse or promote products
14 * derived from this software without specific prior written permission.
15 *
16 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
17 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
18 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
19 * ARE DISCLAIMED. IN NO EVENT SHALL Google LLC BE LIABLE FOR ANY DIRECT,
20 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
21 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
22 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
23 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
25 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26 */
27
28#include "python/map.h"
29
30#include "python/convert.h"
31#include "python/message.h"
32#include "python/protobuf.h"
Protobuf Team Bote4635f22022-06-21 10:43:08 -070033#include "upb/map.h"
Eric Saloedecfd52022-09-15 10:26:14 -070034#include "upb/reflection/def.h"
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080035
36// -----------------------------------------------------------------------------
37// MapContainer
38// -----------------------------------------------------------------------------
39
40typedef struct {
Joshua Haberman1c955f32022-01-12 07:19:28 -080041 PyObject_HEAD;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080042 PyObject* arena;
Joshua Haberman1c955f32022-01-12 07:19:28 -080043 // The field descriptor (upb_FieldDef*).
Joshua Habermanfaac2d82021-12-30 01:04:49 -080044 // The low bit indicates whether the container is reified (see ptr below).
Joshua Haberman58960e02021-12-30 10:50:41 -080045 // - low bit set: repeated field is a stub (empty map, no underlying data).
Joshua Haberman1c955f32022-01-12 07:19:28 -080046 // - low bit clear: repeated field is reified (points to upb_Array).
Joshua Habermanfaac2d82021-12-30 01:04:49 -080047 uintptr_t field;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080048 union {
Joshua Habermanfaac2d82021-12-30 01:04:49 -080049 PyObject* parent; // stub: owning pointer to parent message.
Joshua Haberman1c955f32022-01-12 07:19:28 -080050 upb_Map* map; // reified: the data for this array.
Joshua Habermanfaac2d82021-12-30 01:04:49 -080051 } ptr;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080052 int version;
53} PyUpb_MapContainer;
54
55static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map);
56
Joshua Haberman21d716d2021-12-30 01:07:47 -080057static bool PyUpb_MapContainer_IsStub(PyUpb_MapContainer* self) {
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080058 return self->field & 1;
59}
60
Joshua Haberman58960e02021-12-30 10:50:41 -080061// If the map is reified, returns it. Otherwise, returns NULL.
62// If NULL is returned, the object is empty and has no underlying data.
Joshua Haberman1c955f32022-01-12 07:19:28 -080063static upb_Map* PyUpb_MapContainer_GetIfReified(PyUpb_MapContainer* self) {
Joshua Haberman21d716d2021-12-30 01:07:47 -080064 return PyUpb_MapContainer_IsStub(self) ? NULL : self->ptr.map;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080065}
66
Joshua Haberman1c955f32022-01-12 07:19:28 -080067static const upb_FieldDef* PyUpb_MapContainer_GetField(
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080068 PyUpb_MapContainer* self) {
Joshua Haberman1c955f32022-01-12 07:19:28 -080069 return (const upb_FieldDef*)(self->field & ~(uintptr_t)1);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080070}
71
72static void PyUpb_MapContainer_Dealloc(void* _self) {
73 PyUpb_MapContainer* self = _self;
74 Py_DECREF(self->arena);
Joshua Haberman21d716d2021-12-30 01:07:47 -080075 if (PyUpb_MapContainer_IsStub(self)) {
Joshua Haberman12fffeb2022-05-17 14:00:49 -070076 PyUpb_Message_CacheDelete(self->ptr.parent,
77 PyUpb_MapContainer_GetField(self));
Joshua Habermanfaac2d82021-12-30 01:04:49 -080078 Py_DECREF(self->ptr.parent);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080079 } else {
Joshua Habermanfaac2d82021-12-30 01:04:49 -080080 PyUpb_ObjCache_Delete(self->ptr.map);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080081 }
82 PyUpb_Dealloc(_self);
83}
84
Joshua Haberman1c955f32022-01-12 07:19:28 -080085PyTypeObject* PyUpb_MapContainer_GetClass(const upb_FieldDef* f) {
86 assert(upb_FieldDef_IsMap(f));
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080087 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
Joshua Haberman16daf912022-06-29 08:18:23 -070088 const upb_FieldDef* val =
89 upb_MessageDef_Field(upb_FieldDef_MessageSubDef(f), 1);
90 assert(upb_FieldDef_Number(val) == 2);
91 return upb_FieldDef_IsSubMessage(val) ? state->message_map_container_type
92 : state->scalar_map_container_type;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -080093}
94
Joshua Haberman1c955f32022-01-12 07:19:28 -080095PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_FieldDef* f,
Joshua Habermanfaac2d82021-12-30 01:04:49 -080096 PyObject* arena) {
Joshua Habermane5d8d282022-01-05 15:47:59 -080097 // We only create stubs when the parent is reified, by convention. However
98 // this is not an invariant: the parent could become reified at any time.
Joshua Haberman12fffeb2022-05-17 14:00:49 -070099 assert(PyUpb_Message_GetIfReified(parent) == NULL);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800100 PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800101 PyUpb_MapContainer* map = (void*)PyType_GenericAlloc(cls, 0);
102 map->arena = arena;
103 map->field = (uintptr_t)f | 1;
Joshua Habermanfaac2d82021-12-30 01:04:49 -0800104 map->ptr.parent = parent;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800105 map->version = 0;
106 Py_INCREF(arena);
107 Py_INCREF(parent);
108 return &map->ob_base;
109}
110
Joshua Haberman1c955f32022-01-12 07:19:28 -0800111void PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map) {
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800112 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Haberman1be47592022-01-09 10:27:09 -0800113 if (!map) {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800114 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
115 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
116 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
117 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
118 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
119 map = upb_Map_New(arena, upb_FieldDef_CType(key_f),
120 upb_FieldDef_CType(val_f));
Joshua Haberman1be47592022-01-09 10:27:09 -0800121 }
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800122 PyUpb_ObjCache_Add(map, &self->ob_base);
Joshua Habermanfaac2d82021-12-30 01:04:49 -0800123 Py_DECREF(self->ptr.parent);
124 self->ptr.map = map; // Overwrites self->ptr.parent.
Joshua Haberman1c46e722021-12-30 08:54:42 -0800125 self->field &= ~(uintptr_t)1;
Joshua Haberman21d716d2021-12-30 01:07:47 -0800126 assert(!PyUpb_MapContainer_IsStub(self));
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800127}
128
129void PyUpb_MapContainer_Invalidate(PyObject* obj) {
130 PyUpb_MapContainer* self = (PyUpb_MapContainer*)obj;
131 self->version++;
132}
133
Joshua Haberman1c955f32022-01-12 07:19:28 -0800134upb_Map* PyUpb_MapContainer_EnsureReified(PyObject* _self) {
Joshua Haberman54b775d2022-01-05 12:55:52 -0800135 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800136 self->version++;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800137 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800138 if (map) return map; // Already writable.
139
Joshua Haberman1c955f32022-01-12 07:19:28 -0800140 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
141 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
142 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
143 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
144 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
145 map =
146 upb_Map_New(arena, upb_FieldDef_CType(key_f), upb_FieldDef_CType(val_f));
147 upb_MessageValue msgval = {.map_val = map};
Joshua Haberman12fffeb2022-05-17 14:00:49 -0700148 PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f, msgval);
Joshua Habermanfaac2d82021-12-30 01:04:49 -0800149 PyUpb_MapContainer_Reify((PyObject*)self, map);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800150 return map;
151}
152
Joshua Habermand72d9842022-05-07 13:25:14 -0700153bool PyUpb_MapContainer_Set(PyUpb_MapContainer* self, upb_Map* map,
154 upb_MessageValue key, upb_MessageValue val,
155 upb_Arena* arena) {
156 switch (upb_Map_Insert(map, key, val, arena)) {
157 case kUpb_MapInsertStatus_Inserted:
158 return true;
159 case kUpb_MapInsertStatus_Replaced:
160 // We did not insert a new key, undo the previous invalidate.
161 self->version--;
162 return true;
163 case kUpb_MapInsertStatus_OutOfMemory:
164 return false;
165 }
Joshua Habermane4c7f822022-05-17 10:15:21 -0700166 return false; // Unreachable, silence compiler warning.
Joshua Habermand72d9842022-05-07 13:25:14 -0700167}
168
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800169int PyUpb_MapContainer_AssignSubscript(PyObject* _self, PyObject* key,
170 PyObject* val) {
171 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800172 upb_Map* map = PyUpb_MapContainer_EnsureReified(_self);
173 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
174 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
175 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
176 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
177 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
178 upb_MessageValue u_key, u_val;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800179 if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return -1;
180
181 if (val) {
182 if (!PyUpb_PyToUpb(val, val_f, &u_val, arena)) return -1;
Joshua Habermand72d9842022-05-07 13:25:14 -0700183 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return -1;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800184 } else {
Eric Salo0c6b72d2022-12-29 09:05:43 -0800185 if (!upb_Map_Delete(map, u_key, NULL)) {
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800186 PyErr_Format(PyExc_KeyError, "Key not present in map");
187 return -1;
188 }
189 }
190 return 0;
191}
192
Joshua Haberman1c46e722021-12-30 08:54:42 -0800193PyObject* PyUpb_MapContainer_Subscript(PyObject* _self, PyObject* key) {
194 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800195 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
196 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
197 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
198 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
199 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
200 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
201 upb_MessageValue u_key, u_val;
Joshua Haberman1c46e722021-12-30 08:54:42 -0800202 if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800203 if (!map || !upb_Map_Get(map, u_key, &u_val)) {
Joshua Haberman54b775d2022-01-05 12:55:52 -0800204 map = PyUpb_MapContainer_EnsureReified(_self);
Joshua Haberman1c955f32022-01-12 07:19:28 -0800205 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
206 if (upb_FieldDef_IsSubMessage(val_f)) {
Eric Salo41335a02022-10-10 13:53:47 -0700207 const upb_Message* m = upb_FieldDef_MessageSubDef(val_f);
208 const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
209 u_val.msg_val = upb_Message_New(layout, arena);
Joshua Haberman1c46e722021-12-30 08:54:42 -0800210 } else {
211 memset(&u_val, 0, sizeof(u_val));
212 }
Joshua Habermand72d9842022-05-07 13:25:14 -0700213 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return false;
Joshua Haberman1c46e722021-12-30 08:54:42 -0800214 }
215 return PyUpb_UpbToPy(u_val, val_f, self->arena);
216}
217
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800218PyObject* PyUpb_MapContainer_Contains(PyObject* _self, PyObject* key) {
219 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800220 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800221 if (!map) Py_RETURN_FALSE;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800222 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
223 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
224 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
225 upb_MessageValue u_key;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800226 if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800227 if (upb_Map_Get(map, u_key, NULL)) {
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800228 Py_RETURN_TRUE;
229 } else {
230 Py_RETURN_FALSE;
231 }
232}
233
234PyObject* PyUpb_MapContainer_Clear(PyObject* _self, PyObject* key) {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800235 upb_Map* map = PyUpb_MapContainer_EnsureReified(_self);
236 upb_Map_Clear(map);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800237 Py_RETURN_NONE;
238}
239
240static PyObject* PyUpb_MapContainer_Get(PyObject* _self, PyObject* args,
241 PyObject* kwargs) {
242 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
243 static const char* kwlist[] = {"key", "default", NULL};
244 PyObject* key;
245 PyObject* default_value = NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800246 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800247 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O", (char**)kwlist, &key,
248 &default_value)) {
249 return NULL;
250 }
251
Joshua Haberman1c955f32022-01-12 07:19:28 -0800252 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
253 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
254 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
255 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
256 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
257 upb_MessageValue u_key, u_val;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800258 if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800259 if (map && upb_Map_Get(map, u_key, &u_val)) {
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800260 return PyUpb_UpbToPy(u_val, val_f, self->arena);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800261 }
Joshua Haberman1c46e722021-12-30 08:54:42 -0800262 if (default_value) {
263 Py_INCREF(default_value);
264 return default_value;
265 }
266 Py_RETURN_NONE;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800267}
268
269static PyObject* PyUpb_MapContainer_GetEntryClass(PyObject* _self,
270 PyObject* arg) {
271 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800272 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
273 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800274 return PyUpb_Descriptor_GetClass(entry_m);
275}
276
277Py_ssize_t PyUpb_MapContainer_Length(PyObject* _self) {
278 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800279 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
280 return map ? upb_Map_Size(map) : 0;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800281}
282
283PyUpb_MapContainer* PyUpb_MapContainer_Check(PyObject* _self) {
284 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
285 if (!PyObject_TypeCheck(_self, state->message_map_container_type) &&
286 !PyObject_TypeCheck(_self, state->scalar_map_container_type)) {
287 PyErr_Format(PyExc_TypeError, "Expected protobuf map, but got %R", _self);
288 return NULL;
289 }
290 return (PyUpb_MapContainer*)_self;
291}
292
Joshua Haberman12fffeb2022-05-17 14:00:49 -0700293int PyUpb_Message_InitMapAttributes(PyObject* map, PyObject* value,
294 const upb_FieldDef* f);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800295
296static PyObject* PyUpb_MapContainer_MergeFrom(PyObject* _self, PyObject* _arg) {
297 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800298 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800299
300 if (PyDict_Check(_arg)) {
301 return PyErr_Format(PyExc_AttributeError, "Merging of dict is not allowed");
302 }
303
Joshua Haberman12fffeb2022-05-17 14:00:49 -0700304 if (PyUpb_Message_InitMapAttributes(_self, _arg, f) < 0) {
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800305 return NULL;
306 }
307
308 Py_RETURN_NONE;
309}
310
311static PyObject* PyUpb_MapContainer_Repr(PyObject* _self) {
312 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800313 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800314 PyObject* dict = PyDict_New();
315 if (map) {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800316 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
317 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
318 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
319 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
320 size_t iter = kUpb_Map_Begin;
Eric Salo03b1dee2022-11-21 08:38:37 -0800321 upb_MessageValue map_key, map_val;
322 while (upb_Map_Next(map, &map_key, &map_val, &iter)) {
323 PyObject* key = PyUpb_UpbToPy(map_key, key_f, self->arena);
324 PyObject* val = PyUpb_UpbToPy(map_val, val_f, self->arena);
Joshua Haberman1c46e722021-12-30 08:54:42 -0800325 if (!key || !val) {
326 Py_XDECREF(key);
327 Py_XDECREF(val);
Joshua Haberman58960e02021-12-30 10:50:41 -0800328 Py_DECREF(dict);
Joshua Haberman1c46e722021-12-30 08:54:42 -0800329 return NULL;
330 }
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800331 PyDict_SetItem(dict, key, val);
332 Py_DECREF(key);
333 Py_DECREF(val);
334 }
335 }
336 PyObject* repr = PyObject_Repr(dict);
337 Py_DECREF(dict);
338 return repr;
339}
340
Joshua Haberman1c955f32022-01-12 07:19:28 -0800341PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_Map* map,
342 const upb_FieldDef* f,
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800343 PyObject* arena) {
Joshua Haberman1c46e722021-12-30 08:54:42 -0800344 PyUpb_MapContainer* ret = (void*)PyUpb_ObjCache_Get(map);
345 if (ret) return &ret->ob_base;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800346
Joshua Haberman1c46e722021-12-30 08:54:42 -0800347 PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
348 ret = (void*)PyType_GenericAlloc(cls, 0);
349 ret->arena = arena;
350 ret->field = (uintptr_t)f;
351 ret->ptr.map = map;
352 ret->version = 0;
353 Py_INCREF(arena);
354 PyUpb_ObjCache_Add(map, &ret->ob_base);
355 return &ret->ob_base;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800356}
357
358// -----------------------------------------------------------------------------
359// ScalarMapContainer
360// -----------------------------------------------------------------------------
361
362static PyMethodDef PyUpb_ScalarMapContainer_Methods[] = {
363 {"__contains__", PyUpb_MapContainer_Contains, METH_O,
364 "Tests whether a key is a member of the map."},
365 {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
366 "Removes all elements from the map."},
367 {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
368 "Gets the value for the given key if present, or otherwise a default"},
369 {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
370 "Return the class used to build Entries of (key, value) pairs."},
371 {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
372 "Merges a map into the current map."},
373 /*
374 { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
375 "Makes a deep copy of the class." },
376 { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
377 "Outputs picklable representation of the repeated field." },
378 */
379 {NULL, NULL},
380};
381
382static PyType_Slot PyUpb_ScalarMapContainer_Slots[] = {
383 {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
384 {Py_mp_length, PyUpb_MapContainer_Length},
385 {Py_mp_subscript, PyUpb_MapContainer_Subscript},
386 {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
387 {Py_tp_methods, PyUpb_ScalarMapContainer_Methods},
388 {Py_tp_iter, PyUpb_MapIterator_New},
389 {Py_tp_repr, PyUpb_MapContainer_Repr},
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800390 {0, NULL},
391};
392
393static PyType_Spec PyUpb_ScalarMapContainer_Spec = {
394 PYUPB_MODULE_NAME ".ScalarMapContainer",
395 sizeof(PyUpb_MapContainer),
396 0,
397 Py_TPFLAGS_DEFAULT,
398 PyUpb_ScalarMapContainer_Slots,
399};
400
401// -----------------------------------------------------------------------------
402// MessageMapContainer
403// -----------------------------------------------------------------------------
404
405static PyMethodDef PyUpb_MessageMapContainer_Methods[] = {
406 {"__contains__", PyUpb_MapContainer_Contains, METH_O,
407 "Tests whether the map contains this element."},
408 {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
409 "Removes all elements from the map."},
410 {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
411 "Gets the value for the given key if present, or otherwise a default"},
412 {"get_or_create", PyUpb_MapContainer_Subscript, METH_O,
413 "Alias for getitem, useful to make explicit that the map is mutated."},
414 {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
415 "Return the class used to build Entries of (key, value) pairs."},
416 {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
417 "Merges a map into the current map."},
418 /*
419 { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
420 "Makes a deep copy of the class." },
421 { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
422 "Outputs picklable representation of the repeated field." },
423 */
424 {NULL, NULL},
425};
426
427static PyType_Slot PyUpb_MessageMapContainer_Slots[] = {
428 {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
429 {Py_mp_length, PyUpb_MapContainer_Length},
430 {Py_mp_subscript, PyUpb_MapContainer_Subscript},
431 {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
432 {Py_tp_methods, PyUpb_MessageMapContainer_Methods},
433 {Py_tp_iter, PyUpb_MapIterator_New},
434 {Py_tp_repr, PyUpb_MapContainer_Repr},
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800435 {0, NULL}};
436
437static PyType_Spec PyUpb_MessageMapContainer_Spec = {
438 PYUPB_MODULE_NAME ".MessageMapContainer", sizeof(PyUpb_MapContainer), 0,
439 Py_TPFLAGS_DEFAULT, PyUpb_MessageMapContainer_Slots};
440
441// -----------------------------------------------------------------------------
442// MapIterator
443// -----------------------------------------------------------------------------
444
445typedef struct {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800446 PyObject_HEAD;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800447 PyUpb_MapContainer* map; // We own a reference.
448 size_t iter;
449 int version;
450} PyUpb_MapIterator;
451
452static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map) {
453 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
454 PyUpb_MapIterator* iter =
455 (void*)PyType_GenericAlloc(state->map_iterator_type, 0);
456 iter->map = map;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800457 iter->iter = kUpb_Map_Begin;
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800458 iter->version = map->version;
459 Py_INCREF(map);
460 return &iter->ob_base;
461}
462
463static void PyUpb_MapIterator_Dealloc(void* _self) {
464 PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
465 Py_DECREF(&self->map->ob_base);
466 PyUpb_Dealloc(_self);
467}
468
469PyObject* PyUpb_MapIterator_IterNext(PyObject* _self) {
470 PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
471 if (self->version != self->map->version) {
472 return PyErr_Format(PyExc_RuntimeError, "Map modified during iteration.");
473 }
Joshua Haberman1c955f32022-01-12 07:19:28 -0800474 upb_Map* map = PyUpb_MapContainer_GetIfReified(self->map);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800475 if (!map) return NULL;
Eric Salo03b1dee2022-11-21 08:38:37 -0800476 upb_MessageValue key, val;
477 if (!upb_Map_Next(map, &key, &val, &self->iter)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800478 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self->map);
479 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
480 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
Joshua Habermanbf74b3e2021-12-30 00:52:12 -0800481 return PyUpb_UpbToPy(key, key_f, self->map->arena);
482}
483
484static PyType_Slot PyUpb_MapIterator_Slots[] = {
485 {Py_tp_dealloc, PyUpb_MapIterator_Dealloc},
486 {Py_tp_iter, PyObject_SelfIter},
487 {Py_tp_iternext, PyUpb_MapIterator_IterNext},
488 {0, NULL}};
489
490static PyType_Spec PyUpb_MapIterator_Spec = {
491 PYUPB_MODULE_NAME ".MapIterator", sizeof(PyUpb_MapIterator), 0,
492 Py_TPFLAGS_DEFAULT, PyUpb_MapIterator_Slots};
493
494// -----------------------------------------------------------------------------
495// Top Level
496// -----------------------------------------------------------------------------
497
498static PyObject* GetMutableMappingBase(void) {
499 PyObject* collections = NULL;
500 PyObject* mapping = NULL;
501 PyObject* bases = NULL;
502 if ((collections = PyImport_ImportModule("collections.abc")) &&
503 (mapping = PyObject_GetAttrString(collections, "MutableMapping"))) {
504 bases = Py_BuildValue("(O)", mapping);
505 }
506 Py_XDECREF(collections);
507 Py_XDECREF(mapping);
508 return bases;
509}
510
511bool PyUpb_Map_Init(PyObject* m) {
512 PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
513 PyObject* bases = GetMutableMappingBase();
514 if (!bases) return false;
515
516 state->message_map_container_type =
517 PyUpb_AddClassWithBases(m, &PyUpb_MessageMapContainer_Spec, bases);
518 state->scalar_map_container_type =
519 PyUpb_AddClassWithBases(m, &PyUpb_ScalarMapContainer_Spec, bases);
520 state->map_iterator_type = PyUpb_AddClass(m, &PyUpb_MapIterator_Spec);
521
522 Py_DECREF(bases);
523
524 return state->message_map_container_type &&
525 state->scalar_map_container_type && state->map_iterator_type;
526}