blob: a1d75de9a127069aaca047c11523c2d6481da83c [file] [log] [blame]
Adam Cozzette501ecec2023-09-26 14:36:20 -07001// Protocol Buffers - Google's data interchange format
2// Copyright 2023 Google LLC. All rights reserved.
3// https://developers.google.com/protocol-buffers/
4//
5// Redistribution and use in source and binary forms, with or without
6// modification, are permitted provided that the following conditions are
7// met:
8//
9// * Redistributions of source code must retain the above copyright
10// notice, this list of conditions and the following disclaimer.
11// * Redistributions in binary form must reproduce the above
12// copyright notice, this list of conditions and the following disclaimer
13// in the documentation and/or other materials provided with the
14// distribution.
15// * Neither the name of Google LLC nor the names of its
16// contributors may be used to endorse or promote products derived from
17// this software without specific prior written permission.
18//
19// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
20// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
21// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
22// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
23// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
24// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
25// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
26// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
27// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
28// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30
31#include "python/map.h"
32
33#include "python/convert.h"
34#include "python/message.h"
35#include "python/protobuf.h"
Eric Salo07fba1d2023-09-29 14:50:56 -070036#include "upb/message/map.h"
Adam Cozzette501ecec2023-09-26 14:36:20 -070037#include "upb/reflection/def.h"
38
39// -----------------------------------------------------------------------------
40// MapContainer
41// -----------------------------------------------------------------------------
42
43typedef struct {
44 PyObject_HEAD;
45 PyObject* arena;
46 // The field descriptor (upb_FieldDef*).
47 // The low bit indicates whether the container is reified (see ptr below).
48 // - low bit set: repeated field is a stub (empty map, no underlying data).
49 // - low bit clear: repeated field is reified (points to upb_Array).
50 uintptr_t field;
51 union {
52 PyObject* parent; // stub: owning pointer to parent message.
53 upb_Map* map; // reified: the data for this array.
54 } ptr;
55 int version;
56} PyUpb_MapContainer;
57
58static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map);
59
60static bool PyUpb_MapContainer_IsStub(PyUpb_MapContainer* self) {
61 return self->field & 1;
62}
63
64// If the map is reified, returns it. Otherwise, returns NULL.
65// If NULL is returned, the object is empty and has no underlying data.
66static upb_Map* PyUpb_MapContainer_GetIfReified(PyUpb_MapContainer* self) {
67 return PyUpb_MapContainer_IsStub(self) ? NULL : self->ptr.map;
68}
69
70static const upb_FieldDef* PyUpb_MapContainer_GetField(
71 PyUpb_MapContainer* self) {
72 return (const upb_FieldDef*)(self->field & ~(uintptr_t)1);
73}
74
75static void PyUpb_MapContainer_Dealloc(void* _self) {
76 PyUpb_MapContainer* self = _self;
77 Py_DECREF(self->arena);
78 if (PyUpb_MapContainer_IsStub(self)) {
79 PyUpb_Message_CacheDelete(self->ptr.parent,
80 PyUpb_MapContainer_GetField(self));
81 Py_DECREF(self->ptr.parent);
82 } else {
83 PyUpb_ObjCache_Delete(self->ptr.map);
84 }
85 PyUpb_Dealloc(_self);
86}
87
88PyTypeObject* PyUpb_MapContainer_GetClass(const upb_FieldDef* f) {
89 assert(upb_FieldDef_IsMap(f));
90 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
91 const upb_FieldDef* val =
92 upb_MessageDef_Field(upb_FieldDef_MessageSubDef(f), 1);
93 assert(upb_FieldDef_Number(val) == 2);
94 return upb_FieldDef_IsSubMessage(val) ? state->message_map_container_type
95 : state->scalar_map_container_type;
96}
97
98PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_FieldDef* f,
99 PyObject* arena) {
100 // We only create stubs when the parent is reified, by convention. However
101 // this is not an invariant: the parent could become reified at any time.
102 assert(PyUpb_Message_GetIfReified(parent) == NULL);
103 PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
104 PyUpb_MapContainer* map = (void*)PyType_GenericAlloc(cls, 0);
105 map->arena = arena;
106 map->field = (uintptr_t)f | 1;
107 map->ptr.parent = parent;
108 map->version = 0;
109 Py_INCREF(arena);
110 Py_INCREF(parent);
111 return &map->ob_base;
112}
113
114void PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map) {
115 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
116 if (!map) {
117 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
118 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
119 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
120 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
121 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
122 map = upb_Map_New(arena, upb_FieldDef_CType(key_f),
123 upb_FieldDef_CType(val_f));
124 }
125 PyUpb_ObjCache_Add(map, &self->ob_base);
126 Py_DECREF(self->ptr.parent);
127 self->ptr.map = map; // Overwrites self->ptr.parent.
128 self->field &= ~(uintptr_t)1;
129 assert(!PyUpb_MapContainer_IsStub(self));
130}
131
132void PyUpb_MapContainer_Invalidate(PyObject* obj) {
133 PyUpb_MapContainer* self = (PyUpb_MapContainer*)obj;
134 self->version++;
135}
136
137upb_Map* PyUpb_MapContainer_EnsureReified(PyObject* _self) {
138 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
139 self->version++;
140 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
141 if (map) return map; // Already writable.
142
143 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
144 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
145 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
146 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
147 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
148 map =
149 upb_Map_New(arena, upb_FieldDef_CType(key_f), upb_FieldDef_CType(val_f));
150 upb_MessageValue msgval = {.map_val = map};
151 PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f, msgval);
152 PyUpb_MapContainer_Reify((PyObject*)self, map);
153 return map;
154}
155
156bool PyUpb_MapContainer_Set(PyUpb_MapContainer* self, upb_Map* map,
157 upb_MessageValue key, upb_MessageValue val,
158 upb_Arena* arena) {
159 switch (upb_Map_Insert(map, key, val, arena)) {
160 case kUpb_MapInsertStatus_Inserted:
161 return true;
162 case kUpb_MapInsertStatus_Replaced:
163 // We did not insert a new key, undo the previous invalidate.
164 self->version--;
165 return true;
166 case kUpb_MapInsertStatus_OutOfMemory:
167 return false;
168 }
169 return false; // Unreachable, silence compiler warning.
170}
171
172int PyUpb_MapContainer_AssignSubscript(PyObject* _self, PyObject* key,
173 PyObject* val) {
174 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
175 upb_Map* map = PyUpb_MapContainer_EnsureReified(_self);
176 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
177 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
178 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
179 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
180 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
181 upb_MessageValue u_key, u_val;
182 if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return -1;
183
184 if (val) {
185 if (!PyUpb_PyToUpb(val, val_f, &u_val, arena)) return -1;
186 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return -1;
187 } else {
188 if (!upb_Map_Delete(map, u_key, NULL)) {
189 PyErr_Format(PyExc_KeyError, "Key not present in map");
190 return -1;
191 }
192 }
193 return 0;
194}
195
196PyObject* PyUpb_MapContainer_Subscript(PyObject* _self, PyObject* key) {
197 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
198 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
199 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
200 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
201 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
202 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
203 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
204 upb_MessageValue u_key, u_val;
205 if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;
206 if (!map || !upb_Map_Get(map, u_key, &u_val)) {
207 map = PyUpb_MapContainer_EnsureReified(_self);
208 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
209 if (upb_FieldDef_IsSubMessage(val_f)) {
210 const upb_Message* m = upb_FieldDef_MessageSubDef(val_f);
211 const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
212 u_val.msg_val = upb_Message_New(layout, arena);
213 } else {
214 memset(&u_val, 0, sizeof(u_val));
215 }
216 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return false;
217 }
218 return PyUpb_UpbToPy(u_val, val_f, self->arena);
219}
220
221PyObject* PyUpb_MapContainer_Contains(PyObject* _self, PyObject* key) {
222 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
223 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
224 if (!map) Py_RETURN_FALSE;
225 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
226 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
227 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
228 upb_MessageValue u_key;
229 if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
230 if (upb_Map_Get(map, u_key, NULL)) {
231 Py_RETURN_TRUE;
232 } else {
233 Py_RETURN_FALSE;
234 }
235}
236
237PyObject* PyUpb_MapContainer_Clear(PyObject* _self, PyObject* key) {
238 upb_Map* map = PyUpb_MapContainer_EnsureReified(_self);
239 upb_Map_Clear(map);
240 Py_RETURN_NONE;
241}
242
243static PyObject* PyUpb_MapContainer_Get(PyObject* _self, PyObject* args,
244 PyObject* kwargs) {
245 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
246 static const char* kwlist[] = {"key", "default", NULL};
247 PyObject* key;
248 PyObject* default_value = NULL;
249 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
250 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O", (char**)kwlist, &key,
251 &default_value)) {
252 return NULL;
253 }
254
255 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
256 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
257 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
258 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
259 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
260 upb_MessageValue u_key, u_val;
261 if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;
262 if (map && upb_Map_Get(map, u_key, &u_val)) {
263 return PyUpb_UpbToPy(u_val, val_f, self->arena);
264 }
265 if (default_value) {
266 Py_INCREF(default_value);
267 return default_value;
268 }
269 Py_RETURN_NONE;
270}
271
272static PyObject* PyUpb_MapContainer_GetEntryClass(PyObject* _self,
273 PyObject* arg) {
274 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
275 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
276 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
277 return PyUpb_Descriptor_GetClass(entry_m);
278}
279
280Py_ssize_t PyUpb_MapContainer_Length(PyObject* _self) {
281 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
282 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
283 return map ? upb_Map_Size(map) : 0;
284}
285
286PyUpb_MapContainer* PyUpb_MapContainer_Check(PyObject* _self) {
287 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
288 if (!PyObject_TypeCheck(_self, state->message_map_container_type) &&
289 !PyObject_TypeCheck(_self, state->scalar_map_container_type)) {
290 PyErr_Format(PyExc_TypeError, "Expected protobuf map, but got %R", _self);
291 return NULL;
292 }
293 return (PyUpb_MapContainer*)_self;
294}
295
296int PyUpb_Message_InitMapAttributes(PyObject* map, PyObject* value,
297 const upb_FieldDef* f);
298
299static PyObject* PyUpb_MapContainer_MergeFrom(PyObject* _self, PyObject* _arg) {
300 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
301 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
302
303 if (PyDict_Check(_arg)) {
304 return PyErr_Format(PyExc_AttributeError, "Merging of dict is not allowed");
305 }
306
307 if (PyUpb_Message_InitMapAttributes(_self, _arg, f) < 0) {
308 return NULL;
309 }
310
311 Py_RETURN_NONE;
312}
313
314static PyObject* PyUpb_MapContainer_Repr(PyObject* _self) {
315 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
316 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
317 PyObject* dict = PyDict_New();
318 if (map) {
319 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
320 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
321 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
322 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
323 size_t iter = kUpb_Map_Begin;
324 upb_MessageValue map_key, map_val;
325 while (upb_Map_Next(map, &map_key, &map_val, &iter)) {
326 PyObject* key = PyUpb_UpbToPy(map_key, key_f, self->arena);
327 PyObject* val = PyUpb_UpbToPy(map_val, val_f, self->arena);
328 if (!key || !val) {
329 Py_XDECREF(key);
330 Py_XDECREF(val);
331 Py_DECREF(dict);
332 return NULL;
333 }
334 PyDict_SetItem(dict, key, val);
335 Py_DECREF(key);
336 Py_DECREF(val);
337 }
338 }
339 PyObject* repr = PyObject_Repr(dict);
340 Py_DECREF(dict);
341 return repr;
342}
343
344PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_Map* map,
345 const upb_FieldDef* f,
346 PyObject* arena) {
347 PyUpb_MapContainer* ret = (void*)PyUpb_ObjCache_Get(map);
348 if (ret) return &ret->ob_base;
349
350 PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
351 ret = (void*)PyType_GenericAlloc(cls, 0);
352 ret->arena = arena;
353 ret->field = (uintptr_t)f;
354 ret->ptr.map = map;
355 ret->version = 0;
356 Py_INCREF(arena);
357 PyUpb_ObjCache_Add(map, &ret->ob_base);
358 return &ret->ob_base;
359}
360
361// -----------------------------------------------------------------------------
362// ScalarMapContainer
363// -----------------------------------------------------------------------------
364
365static PyMethodDef PyUpb_ScalarMapContainer_Methods[] = {
366 {"__contains__", PyUpb_MapContainer_Contains, METH_O,
367 "Tests whether a key is a member of the map."},
368 {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
369 "Removes all elements from the map."},
370 {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
371 "Gets the value for the given key if present, or otherwise a default"},
372 {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
373 "Return the class used to build Entries of (key, value) pairs."},
374 {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
375 "Merges a map into the current map."},
376 /*
377 { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
378 "Makes a deep copy of the class." },
379 { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
380 "Outputs picklable representation of the repeated field." },
381 */
382 {NULL, NULL},
383};
384
385static PyType_Slot PyUpb_ScalarMapContainer_Slots[] = {
386 {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
387 {Py_mp_length, PyUpb_MapContainer_Length},
388 {Py_mp_subscript, PyUpb_MapContainer_Subscript},
389 {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
390 {Py_tp_methods, PyUpb_ScalarMapContainer_Methods},
391 {Py_tp_iter, PyUpb_MapIterator_New},
392 {Py_tp_repr, PyUpb_MapContainer_Repr},
393 {0, NULL},
394};
395
396static PyType_Spec PyUpb_ScalarMapContainer_Spec = {
397 PYUPB_MODULE_NAME ".ScalarMapContainer",
398 sizeof(PyUpb_MapContainer),
399 0,
400 Py_TPFLAGS_DEFAULT,
401 PyUpb_ScalarMapContainer_Slots,
402};
403
404// -----------------------------------------------------------------------------
405// MessageMapContainer
406// -----------------------------------------------------------------------------
407
408static PyMethodDef PyUpb_MessageMapContainer_Methods[] = {
409 {"__contains__", PyUpb_MapContainer_Contains, METH_O,
410 "Tests whether the map contains this element."},
411 {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
412 "Removes all elements from the map."},
413 {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
414 "Gets the value for the given key if present, or otherwise a default"},
415 {"get_or_create", PyUpb_MapContainer_Subscript, METH_O,
416 "Alias for getitem, useful to make explicit that the map is mutated."},
417 {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
418 "Return the class used to build Entries of (key, value) pairs."},
419 {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
420 "Merges a map into the current map."},
421 /*
422 { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
423 "Makes a deep copy of the class." },
424 { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
425 "Outputs picklable representation of the repeated field." },
426 */
427 {NULL, NULL},
428};
429
430static PyType_Slot PyUpb_MessageMapContainer_Slots[] = {
431 {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
432 {Py_mp_length, PyUpb_MapContainer_Length},
433 {Py_mp_subscript, PyUpb_MapContainer_Subscript},
434 {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
435 {Py_tp_methods, PyUpb_MessageMapContainer_Methods},
436 {Py_tp_iter, PyUpb_MapIterator_New},
437 {Py_tp_repr, PyUpb_MapContainer_Repr},
438 {0, NULL}};
439
440static PyType_Spec PyUpb_MessageMapContainer_Spec = {
441 PYUPB_MODULE_NAME ".MessageMapContainer", sizeof(PyUpb_MapContainer), 0,
442 Py_TPFLAGS_DEFAULT, PyUpb_MessageMapContainer_Slots};
443
444// -----------------------------------------------------------------------------
445// MapIterator
446// -----------------------------------------------------------------------------
447
448typedef struct {
449 PyObject_HEAD;
450 PyUpb_MapContainer* map; // We own a reference.
451 size_t iter;
452 int version;
453} PyUpb_MapIterator;
454
455static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map) {
456 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
457 PyUpb_MapIterator* iter =
458 (void*)PyType_GenericAlloc(state->map_iterator_type, 0);
459 iter->map = map;
460 iter->iter = kUpb_Map_Begin;
461 iter->version = map->version;
462 Py_INCREF(map);
463 return &iter->ob_base;
464}
465
466static void PyUpb_MapIterator_Dealloc(void* _self) {
467 PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
468 Py_DECREF(&self->map->ob_base);
469 PyUpb_Dealloc(_self);
470}
471
472PyObject* PyUpb_MapIterator_IterNext(PyObject* _self) {
473 PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
474 if (self->version != self->map->version) {
475 return PyErr_Format(PyExc_RuntimeError, "Map modified during iteration.");
476 }
477 upb_Map* map = PyUpb_MapContainer_GetIfReified(self->map);
478 if (!map) return NULL;
479 upb_MessageValue key, val;
480 if (!upb_Map_Next(map, &key, &val, &self->iter)) return NULL;
481 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self->map);
482 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
483 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
484 return PyUpb_UpbToPy(key, key_f, self->map->arena);
485}
486
487static PyType_Slot PyUpb_MapIterator_Slots[] = {
488 {Py_tp_dealloc, PyUpb_MapIterator_Dealloc},
489 {Py_tp_iter, PyObject_SelfIter},
490 {Py_tp_iternext, PyUpb_MapIterator_IterNext},
491 {0, NULL}};
492
493static PyType_Spec PyUpb_MapIterator_Spec = {
494 PYUPB_MODULE_NAME ".MapIterator", sizeof(PyUpb_MapIterator), 0,
495 Py_TPFLAGS_DEFAULT, PyUpb_MapIterator_Slots};
496
497// -----------------------------------------------------------------------------
498// Top Level
499// -----------------------------------------------------------------------------
500
501static PyObject* GetMutableMappingBase(void) {
502 PyObject* collections = NULL;
503 PyObject* mapping = NULL;
504 PyObject* bases = NULL;
505 if ((collections = PyImport_ImportModule("collections.abc")) &&
506 (mapping = PyObject_GetAttrString(collections, "MutableMapping"))) {
507 bases = Py_BuildValue("(O)", mapping);
508 }
509 Py_XDECREF(collections);
510 Py_XDECREF(mapping);
511 return bases;
512}
513
514bool PyUpb_Map_Init(PyObject* m) {
515 PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
516 PyObject* bases = GetMutableMappingBase();
517 if (!bases) return false;
518
519 state->message_map_container_type =
520 PyUpb_AddClassWithBases(m, &PyUpb_MessageMapContainer_Spec, bases);
521 state->scalar_map_container_type =
522 PyUpb_AddClassWithBases(m, &PyUpb_ScalarMapContainer_Spec, bases);
523 state->map_iterator_type = PyUpb_AddClass(m, &PyUpb_MapIterator_Spec);
524
525 Py_DECREF(bases);
526
527 return state->message_map_container_type &&
528 state->scalar_map_container_type && state->map_iterator_type;
529}