1 /*
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"
33 #include "upb/map.h"
34 #include "upb/reflection/def.h"
35
36 // -----------------------------------------------------------------------------
37 // MapContainer
38 // -----------------------------------------------------------------------------
39
40 typedef struct {
41 PyObject_HEAD;
42 PyObject* arena;
43 // The field descriptor (upb_FieldDef*).
44 // The low bit indicates whether the container is reified (see ptr below).
45 // - low bit set: repeated field is a stub (empty map, no underlying data).
46 // - low bit clear: repeated field is reified (points to upb_Array).
47 uintptr_t field;
48 union {
49 PyObject* parent; // stub: owning pointer to parent message.
50 upb_Map* map; // reified: the data for this array.
51 } ptr;
52 int version;
53 } PyUpb_MapContainer;
54
55 static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map);
56
PyUpb_MapContainer_IsStub(PyUpb_MapContainer * self)57 static bool PyUpb_MapContainer_IsStub(PyUpb_MapContainer* self) {
58 return self->field & 1;
59 }
60
61 // If the map is reified, returns it. Otherwise, returns NULL.
62 // If NULL is returned, the object is empty and has no underlying data.
PyUpb_MapContainer_GetIfReified(PyUpb_MapContainer * self)63 static upb_Map* PyUpb_MapContainer_GetIfReified(PyUpb_MapContainer* self) {
64 return PyUpb_MapContainer_IsStub(self) ? NULL : self->ptr.map;
65 }
66
PyUpb_MapContainer_GetField(PyUpb_MapContainer * self)67 static const upb_FieldDef* PyUpb_MapContainer_GetField(
68 PyUpb_MapContainer* self) {
69 return (const upb_FieldDef*)(self->field & ~(uintptr_t)1);
70 }
71
PyUpb_MapContainer_Dealloc(void * _self)72 static void PyUpb_MapContainer_Dealloc(void* _self) {
73 PyUpb_MapContainer* self = _self;
74 Py_DECREF(self->arena);
75 if (PyUpb_MapContainer_IsStub(self)) {
76 PyUpb_Message_CacheDelete(self->ptr.parent,
77 PyUpb_MapContainer_GetField(self));
78 Py_DECREF(self->ptr.parent);
79 } else {
80 PyUpb_ObjCache_Delete(self->ptr.map);
81 }
82 PyUpb_Dealloc(_self);
83 }
84
PyUpb_MapContainer_GetClass(const upb_FieldDef * f)85 PyTypeObject* PyUpb_MapContainer_GetClass(const upb_FieldDef* f) {
86 assert(upb_FieldDef_IsMap(f));
87 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
88 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;
93 }
94
PyUpb_MapContainer_NewStub(PyObject * parent,const upb_FieldDef * f,PyObject * arena)95 PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_FieldDef* f,
96 PyObject* arena) {
97 // 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.
99 assert(PyUpb_Message_GetIfReified(parent) == NULL);
100 PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
101 PyUpb_MapContainer* map = (void*)PyType_GenericAlloc(cls, 0);
102 map->arena = arena;
103 map->field = (uintptr_t)f | 1;
104 map->ptr.parent = parent;
105 map->version = 0;
106 Py_INCREF(arena);
107 Py_INCREF(parent);
108 return &map->ob_base;
109 }
110
PyUpb_MapContainer_Reify(PyObject * _self,upb_Map * map)111 void PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map) {
112 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
113 if (!map) {
114 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));
121 }
122 PyUpb_ObjCache_Add(map, &self->ob_base);
123 Py_DECREF(self->ptr.parent);
124 self->ptr.map = map; // Overwrites self->ptr.parent.
125 self->field &= ~(uintptr_t)1;
126 assert(!PyUpb_MapContainer_IsStub(self));
127 }
128
PyUpb_MapContainer_Invalidate(PyObject * obj)129 void PyUpb_MapContainer_Invalidate(PyObject* obj) {
130 PyUpb_MapContainer* self = (PyUpb_MapContainer*)obj;
131 self->version++;
132 }
133
PyUpb_MapContainer_EnsureReified(PyObject * _self)134 upb_Map* PyUpb_MapContainer_EnsureReified(PyObject* _self) {
135 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
136 self->version++;
137 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
138 if (map) return map; // Already writable.
139
140 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};
148 PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f, msgval);
149 PyUpb_MapContainer_Reify((PyObject*)self, map);
150 return map;
151 }
152
PyUpb_MapContainer_Set(PyUpb_MapContainer * self,upb_Map * map,upb_MessageValue key,upb_MessageValue val,upb_Arena * arena)153 bool 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 }
166 return false; // Unreachable, silence compiler warning.
167 }
168
PyUpb_MapContainer_AssignSubscript(PyObject * _self,PyObject * key,PyObject * val)169 int PyUpb_MapContainer_AssignSubscript(PyObject* _self, PyObject* key,
170 PyObject* val) {
171 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
172 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;
179 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;
183 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return -1;
184 } else {
185 if (!upb_Map_Delete(map, u_key, NULL)) {
186 PyErr_Format(PyExc_KeyError, "Key not present in map");
187 return -1;
188 }
189 }
190 return 0;
191 }
192
PyUpb_MapContainer_Subscript(PyObject * _self,PyObject * key)193 PyObject* PyUpb_MapContainer_Subscript(PyObject* _self, PyObject* key) {
194 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
195 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;
202 if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;
203 if (!map || !upb_Map_Get(map, u_key, &u_val)) {
204 map = PyUpb_MapContainer_EnsureReified(_self);
205 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
206 if (upb_FieldDef_IsSubMessage(val_f)) {
207 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);
210 } else {
211 memset(&u_val, 0, sizeof(u_val));
212 }
213 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return false;
214 }
215 return PyUpb_UpbToPy(u_val, val_f, self->arena);
216 }
217
PyUpb_MapContainer_Contains(PyObject * _self,PyObject * key)218 PyObject* PyUpb_MapContainer_Contains(PyObject* _self, PyObject* key) {
219 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
220 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
221 if (!map) Py_RETURN_FALSE;
222 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;
226 if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
227 if (upb_Map_Get(map, u_key, NULL)) {
228 Py_RETURN_TRUE;
229 } else {
230 Py_RETURN_FALSE;
231 }
232 }
233
PyUpb_MapContainer_Clear(PyObject * _self,PyObject * key)234 PyObject* PyUpb_MapContainer_Clear(PyObject* _self, PyObject* key) {
235 upb_Map* map = PyUpb_MapContainer_EnsureReified(_self);
236 upb_Map_Clear(map);
237 Py_RETURN_NONE;
238 }
239
PyUpb_MapContainer_Get(PyObject * _self,PyObject * args,PyObject * kwargs)240 static 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;
246 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
247 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O", (char**)kwlist, &key,
248 &default_value)) {
249 return NULL;
250 }
251
252 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;
258 if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;
259 if (map && upb_Map_Get(map, u_key, &u_val)) {
260 return PyUpb_UpbToPy(u_val, val_f, self->arena);
261 }
262 if (default_value) {
263 Py_INCREF(default_value);
264 return default_value;
265 }
266 Py_RETURN_NONE;
267 }
268
PyUpb_MapContainer_GetEntryClass(PyObject * _self,PyObject * arg)269 static PyObject* PyUpb_MapContainer_GetEntryClass(PyObject* _self,
270 PyObject* arg) {
271 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
272 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
273 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
274 return PyUpb_Descriptor_GetClass(entry_m);
275 }
276
PyUpb_MapContainer_Length(PyObject * _self)277 Py_ssize_t PyUpb_MapContainer_Length(PyObject* _self) {
278 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
279 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
280 return map ? upb_Map_Size(map) : 0;
281 }
282
PyUpb_MapContainer_Check(PyObject * _self)283 PyUpb_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
293 int PyUpb_Message_InitMapAttributes(PyObject* map, PyObject* value,
294 const upb_FieldDef* f);
295
PyUpb_MapContainer_MergeFrom(PyObject * _self,PyObject * _arg)296 static PyObject* PyUpb_MapContainer_MergeFrom(PyObject* _self, PyObject* _arg) {
297 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
298 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
299
300 if (PyDict_Check(_arg)) {
301 return PyErr_Format(PyExc_AttributeError, "Merging of dict is not allowed");
302 }
303
304 if (PyUpb_Message_InitMapAttributes(_self, _arg, f) < 0) {
305 return NULL;
306 }
307
308 Py_RETURN_NONE;
309 }
310
PyUpb_MapContainer_Repr(PyObject * _self)311 static PyObject* PyUpb_MapContainer_Repr(PyObject* _self) {
312 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
313 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
314 PyObject* dict = PyDict_New();
315 if (map) {
316 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;
321 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);
325 if (!key || !val) {
326 Py_XDECREF(key);
327 Py_XDECREF(val);
328 Py_DECREF(dict);
329 return NULL;
330 }
331 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
PyUpb_MapContainer_GetOrCreateWrapper(upb_Map * map,const upb_FieldDef * f,PyObject * arena)341 PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_Map* map,
342 const upb_FieldDef* f,
343 PyObject* arena) {
344 PyUpb_MapContainer* ret = (void*)PyUpb_ObjCache_Get(map);
345 if (ret) return &ret->ob_base;
346
347 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;
356 }
357
358 // -----------------------------------------------------------------------------
359 // ScalarMapContainer
360 // -----------------------------------------------------------------------------
361
362 static 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
382 static 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},
390 {0, NULL},
391 };
392
393 static 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
405 static 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
427 static 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},
435 {0, NULL}};
436
437 static 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
445 typedef struct {
446 PyObject_HEAD;
447 PyUpb_MapContainer* map; // We own a reference.
448 size_t iter;
449 int version;
450 } PyUpb_MapIterator;
451
PyUpb_MapIterator_New(PyUpb_MapContainer * map)452 static 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;
457 iter->iter = kUpb_Map_Begin;
458 iter->version = map->version;
459 Py_INCREF(map);
460 return &iter->ob_base;
461 }
462
PyUpb_MapIterator_Dealloc(void * _self)463 static 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
PyUpb_MapIterator_IterNext(PyObject * _self)469 PyObject* 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 }
474 upb_Map* map = PyUpb_MapContainer_GetIfReified(self->map);
475 if (!map) return NULL;
476 upb_MessageValue key, val;
477 if (!upb_Map_Next(map, &key, &val, &self->iter)) return NULL;
478 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);
481 return PyUpb_UpbToPy(key, key_f, self->map->arena);
482 }
483
484 static 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
490 static 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
GetMutableMappingBase(void)498 static 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
PyUpb_Map_Init(PyObject * m)511 bool 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 }
527