blob: 4fe7c1911451d5a16e58385bdcf2cb33e634b6da [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.
Adam Cozzette501ecec2023-09-26 14:36:20 -07003//
Protobuf Team Bot0fab7732023-11-20 13:38:15 -08004// Use of this source code is governed by a BSD-style
5// license that can be found in the LICENSE file or at
6// https://developers.google.com/open-source/licenses/bsd
Adam Cozzette501ecec2023-09-26 14:36:20 -07007
8#include "python/map.h"
9
10#include "python/convert.h"
11#include "python/message.h"
12#include "python/protobuf.h"
Eric Salo07fba1d2023-09-29 14:50:56 -070013#include "upb/message/map.h"
Adam Cozzette501ecec2023-09-26 14:36:20 -070014#include "upb/reflection/def.h"
15
16// -----------------------------------------------------------------------------
17// MapContainer
18// -----------------------------------------------------------------------------
19
20typedef struct {
21 PyObject_HEAD;
22 PyObject* arena;
23 // The field descriptor (upb_FieldDef*).
24 // The low bit indicates whether the container is reified (see ptr below).
25 // - low bit set: repeated field is a stub (empty map, no underlying data).
26 // - low bit clear: repeated field is reified (points to upb_Array).
27 uintptr_t field;
28 union {
29 PyObject* parent; // stub: owning pointer to parent message.
30 upb_Map* map; // reified: the data for this array.
31 } ptr;
32 int version;
33} PyUpb_MapContainer;
34
35static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map);
36
37static bool PyUpb_MapContainer_IsStub(PyUpb_MapContainer* self) {
38 return self->field & 1;
39}
40
41// If the map is reified, returns it. Otherwise, returns NULL.
42// If NULL is returned, the object is empty and has no underlying data.
43static upb_Map* PyUpb_MapContainer_GetIfReified(PyUpb_MapContainer* self) {
44 return PyUpb_MapContainer_IsStub(self) ? NULL : self->ptr.map;
45}
46
47static const upb_FieldDef* PyUpb_MapContainer_GetField(
48 PyUpb_MapContainer* self) {
49 return (const upb_FieldDef*)(self->field & ~(uintptr_t)1);
50}
51
52static void PyUpb_MapContainer_Dealloc(void* _self) {
53 PyUpb_MapContainer* self = _self;
54 Py_DECREF(self->arena);
55 if (PyUpb_MapContainer_IsStub(self)) {
56 PyUpb_Message_CacheDelete(self->ptr.parent,
57 PyUpb_MapContainer_GetField(self));
58 Py_DECREF(self->ptr.parent);
59 } else {
60 PyUpb_ObjCache_Delete(self->ptr.map);
61 }
62 PyUpb_Dealloc(_self);
63}
64
Eric Salo33cb42e2023-12-19 17:48:33 -080065static PyTypeObject* PyUpb_MapContainer_GetClass(const upb_FieldDef* f) {
Adam Cozzette501ecec2023-09-26 14:36:20 -070066 assert(upb_FieldDef_IsMap(f));
67 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
68 const upb_FieldDef* val =
69 upb_MessageDef_Field(upb_FieldDef_MessageSubDef(f), 1);
70 assert(upb_FieldDef_Number(val) == 2);
71 return upb_FieldDef_IsSubMessage(val) ? state->message_map_container_type
72 : state->scalar_map_container_type;
73}
74
75PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_FieldDef* f,
76 PyObject* arena) {
77 // We only create stubs when the parent is reified, by convention. However
78 // this is not an invariant: the parent could become reified at any time.
79 assert(PyUpb_Message_GetIfReified(parent) == NULL);
80 PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
81 PyUpb_MapContainer* map = (void*)PyType_GenericAlloc(cls, 0);
82 map->arena = arena;
83 map->field = (uintptr_t)f | 1;
84 map->ptr.parent = parent;
85 map->version = 0;
86 Py_INCREF(arena);
87 Py_INCREF(parent);
88 return &map->ob_base;
89}
90
91void PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map) {
92 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
93 if (!map) {
94 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
95 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
96 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
97 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
98 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
99 map = upb_Map_New(arena, upb_FieldDef_CType(key_f),
100 upb_FieldDef_CType(val_f));
101 }
102 PyUpb_ObjCache_Add(map, &self->ob_base);
103 Py_DECREF(self->ptr.parent);
104 self->ptr.map = map; // Overwrites self->ptr.parent.
105 self->field &= ~(uintptr_t)1;
106 assert(!PyUpb_MapContainer_IsStub(self));
107}
108
109void PyUpb_MapContainer_Invalidate(PyObject* obj) {
110 PyUpb_MapContainer* self = (PyUpb_MapContainer*)obj;
111 self->version++;
112}
113
114upb_Map* PyUpb_MapContainer_EnsureReified(PyObject* _self) {
115 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
116 self->version++;
117 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
118 if (map) return map; // Already writable.
119
120 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
121 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
122 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
123 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
124 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
125 map =
126 upb_Map_New(arena, upb_FieldDef_CType(key_f), upb_FieldDef_CType(val_f));
127 upb_MessageValue msgval = {.map_val = map};
128 PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f, msgval);
129 PyUpb_MapContainer_Reify((PyObject*)self, map);
130 return map;
131}
132
Eric Salo33cb42e2023-12-19 17:48:33 -0800133static bool PyUpb_MapContainer_Set(PyUpb_MapContainer* self, upb_Map* map,
134 upb_MessageValue key, upb_MessageValue val,
135 upb_Arena* arena) {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700136 switch (upb_Map_Insert(map, key, val, arena)) {
137 case kUpb_MapInsertStatus_Inserted:
138 return true;
139 case kUpb_MapInsertStatus_Replaced:
140 // We did not insert a new key, undo the previous invalidate.
141 self->version--;
142 return true;
143 case kUpb_MapInsertStatus_OutOfMemory:
144 return false;
145 }
146 return false; // Unreachable, silence compiler warning.
147}
148
Eric Salo33cb42e2023-12-19 17:48:33 -0800149// Assigns `self[key] = val` for the map `self`.
150static int PyUpb_MapContainer_AssignSubscript(PyObject* _self, PyObject* key,
151 PyObject* val) {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700152 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
153 upb_Map* map = PyUpb_MapContainer_EnsureReified(_self);
154 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
155 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
156 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
157 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
158 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
159 upb_MessageValue u_key, u_val;
Joshua Habermanb0eeb352023-11-01 10:33:11 -0700160 if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return -1;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700161
162 if (val) {
163 if (!PyUpb_PyToUpb(val, val_f, &u_val, arena)) return -1;
164 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return -1;
165 } else {
166 if (!upb_Map_Delete(map, u_key, NULL)) {
167 PyErr_Format(PyExc_KeyError, "Key not present in map");
168 return -1;
169 }
170 }
171 return 0;
172}
173
Eric Salo33cb42e2023-12-19 17:48:33 -0800174static PyObject* PyUpb_MapContainer_Subscript(PyObject* _self, PyObject* key) {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700175 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
176 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
177 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
178 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
179 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
180 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
Adam Cozzette501ecec2023-09-26 14:36:20 -0700181 upb_MessageValue u_key, u_val;
Joshua Habermanb0eeb352023-11-01 10:33:11 -0700182 if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700183 if (!map || !upb_Map_Get(map, u_key, &u_val)) {
184 map = PyUpb_MapContainer_EnsureReified(_self);
185 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
186 if (upb_FieldDef_IsSubMessage(val_f)) {
Eric Salob997cb62023-12-21 08:09:54 -0800187 const upb_MessageDef* m = upb_FieldDef_MessageSubDef(val_f);
Adam Cozzette501ecec2023-09-26 14:36:20 -0700188 const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
189 u_val.msg_val = upb_Message_New(layout, arena);
190 } else {
191 memset(&u_val, 0, sizeof(u_val));
192 }
193 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return false;
194 }
195 return PyUpb_UpbToPy(u_val, val_f, self->arena);
196}
197
Sandy Zhang5b329362024-03-05 17:54:43 -0800198static int PyUpb_MapContainer_Contains(PyObject* _self, PyObject* key) {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700199 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
200 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
Sandy Zhang5b329362024-03-05 17:54:43 -0800201 if (!map) return 0;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700202 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
203 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
204 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
205 upb_MessageValue u_key;
Sandy Zhang5b329362024-03-05 17:54:43 -0800206 if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return -1;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700207 if (upb_Map_Get(map, u_key, NULL)) {
Sandy Zhang5b329362024-03-05 17:54:43 -0800208 return 1;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700209 } else {
Sandy Zhang5b329362024-03-05 17:54:43 -0800210 return 0;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700211 }
212}
213
Eric Salo33cb42e2023-12-19 17:48:33 -0800214static PyObject* PyUpb_MapContainer_Clear(PyObject* _self, PyObject* key) {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700215 upb_Map* map = PyUpb_MapContainer_EnsureReified(_self);
216 upb_Map_Clear(map);
217 Py_RETURN_NONE;
218}
219
Jie Luo81da6b92024-11-12 09:41:02 -0800220static PyObject* PyUpb_ScalarMapContainer_Setdefault(PyObject* _self,
221 PyObject* args) {
222 PyObject* key;
223 PyObject* default_value = Py_None;
224
225 if (!PyArg_UnpackTuple(args, "setdefault", 1, 2, &key, &default_value)) {
226 return NULL;
227 }
228
229 if (default_value == Py_None) {
230 PyErr_Format(PyExc_ValueError,
231 "The value for scalar map setdefault must be set.");
232 return NULL;
233 }
234
235 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
236 upb_Map* map = PyUpb_MapContainer_EnsureReified(_self);
237 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
238 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
239 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
240 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
241 upb_MessageValue u_key, u_val;
242 if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
243 if (upb_Map_Get(map, u_key, &u_val)) {
244 return PyUpb_UpbToPy(u_val, val_f, self->arena);
245 }
246
247 upb_Arena* arena = PyUpb_Arena_Get(self->arena);
248 if (!PyUpb_PyToUpb(default_value, val_f, &u_val, arena)) return NULL;
249 if (!PyUpb_MapContainer_Set(self, map, u_key, u_val, arena)) return NULL;
250
251 Py_INCREF(default_value);
252 return default_value;
253}
254
255static PyObject* PyUpb_MessageMapContainer_Setdefault(PyObject* self,
256 PyObject* args) {
257 PyErr_Format(PyExc_NotImplementedError,
258 "Set message map value directly is not supported.");
259 return NULL;
260}
261
Adam Cozzette501ecec2023-09-26 14:36:20 -0700262static PyObject* PyUpb_MapContainer_Get(PyObject* _self, PyObject* args,
263 PyObject* kwargs) {
264 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
265 static const char* kwlist[] = {"key", "default", NULL};
266 PyObject* key;
267 PyObject* default_value = NULL;
268 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
269 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O", (char**)kwlist, &key,
270 &default_value)) {
271 return NULL;
272 }
273
274 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
275 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
276 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
277 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
Adam Cozzette501ecec2023-09-26 14:36:20 -0700278 upb_MessageValue u_key, u_val;
Joshua Habermanb0eeb352023-11-01 10:33:11 -0700279 if (!PyUpb_PyToUpb(key, key_f, &u_key, NULL)) return NULL;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700280 if (map && upb_Map_Get(map, u_key, &u_val)) {
281 return PyUpb_UpbToPy(u_val, val_f, self->arena);
282 }
283 if (default_value) {
284 Py_INCREF(default_value);
285 return default_value;
286 }
287 Py_RETURN_NONE;
288}
289
290static PyObject* PyUpb_MapContainer_GetEntryClass(PyObject* _self,
291 PyObject* arg) {
292 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
293 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
294 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
295 return PyUpb_Descriptor_GetClass(entry_m);
296}
297
Eric Salo33cb42e2023-12-19 17:48:33 -0800298static Py_ssize_t PyUpb_MapContainer_Length(PyObject* _self) {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700299 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
300 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
301 return map ? upb_Map_Size(map) : 0;
302}
303
Adam Cozzette501ecec2023-09-26 14:36:20 -0700304int PyUpb_Message_InitMapAttributes(PyObject* map, PyObject* value,
305 const upb_FieldDef* f);
306
307static PyObject* PyUpb_MapContainer_MergeFrom(PyObject* _self, PyObject* _arg) {
308 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
309 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
310
311 if (PyDict_Check(_arg)) {
312 return PyErr_Format(PyExc_AttributeError, "Merging of dict is not allowed");
313 }
314
315 if (PyUpb_Message_InitMapAttributes(_self, _arg, f) < 0) {
316 return NULL;
317 }
318
319 Py_RETURN_NONE;
320}
321
322static PyObject* PyUpb_MapContainer_Repr(PyObject* _self) {
323 PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
324 upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
325 PyObject* dict = PyDict_New();
326 if (map) {
327 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
328 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
329 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
330 const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
331 size_t iter = kUpb_Map_Begin;
332 upb_MessageValue map_key, map_val;
333 while (upb_Map_Next(map, &map_key, &map_val, &iter)) {
334 PyObject* key = PyUpb_UpbToPy(map_key, key_f, self->arena);
335 PyObject* val = PyUpb_UpbToPy(map_val, val_f, self->arena);
336 if (!key || !val) {
337 Py_XDECREF(key);
338 Py_XDECREF(val);
339 Py_DECREF(dict);
340 return NULL;
341 }
342 PyDict_SetItem(dict, key, val);
343 Py_DECREF(key);
344 Py_DECREF(val);
345 }
346 }
347 PyObject* repr = PyObject_Repr(dict);
348 Py_DECREF(dict);
349 return repr;
350}
351
352PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_Map* map,
353 const upb_FieldDef* f,
354 PyObject* arena) {
355 PyUpb_MapContainer* ret = (void*)PyUpb_ObjCache_Get(map);
356 if (ret) return &ret->ob_base;
357
358 PyTypeObject* cls = PyUpb_MapContainer_GetClass(f);
359 ret = (void*)PyType_GenericAlloc(cls, 0);
360 ret->arena = arena;
361 ret->field = (uintptr_t)f;
362 ret->ptr.map = map;
363 ret->version = 0;
364 Py_INCREF(arena);
365 PyUpb_ObjCache_Add(map, &ret->ob_base);
366 return &ret->ob_base;
367}
368
369// -----------------------------------------------------------------------------
370// ScalarMapContainer
371// -----------------------------------------------------------------------------
372
373static PyMethodDef PyUpb_ScalarMapContainer_Methods[] = {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700374 {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
375 "Removes all elements from the map."},
Jie Luo81da6b92024-11-12 09:41:02 -0800376 {"setdefault", (PyCFunction)PyUpb_ScalarMapContainer_Setdefault,
377 METH_VARARGS,
378 "If the key does not exist, insert the key, with the specified value"},
Adam Cozzette501ecec2023-09-26 14:36:20 -0700379 {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
380 "Gets the value for the given key if present, or otherwise a default"},
381 {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
382 "Return the class used to build Entries of (key, value) pairs."},
383 {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
384 "Merges a map into the current map."},
385 /*
386 { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
387 "Makes a deep copy of the class." },
388 { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
389 "Outputs picklable representation of the repeated field." },
390 */
391 {NULL, NULL},
392};
393
394static PyType_Slot PyUpb_ScalarMapContainer_Slots[] = {
395 {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
396 {Py_mp_length, PyUpb_MapContainer_Length},
397 {Py_mp_subscript, PyUpb_MapContainer_Subscript},
398 {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
Sandy Zhang5b329362024-03-05 17:54:43 -0800399 {Py_sq_contains, PyUpb_MapContainer_Contains},
Adam Cozzette501ecec2023-09-26 14:36:20 -0700400 {Py_tp_methods, PyUpb_ScalarMapContainer_Methods},
401 {Py_tp_iter, PyUpb_MapIterator_New},
402 {Py_tp_repr, PyUpb_MapContainer_Repr},
403 {0, NULL},
404};
405
406static PyType_Spec PyUpb_ScalarMapContainer_Spec = {
407 PYUPB_MODULE_NAME ".ScalarMapContainer",
408 sizeof(PyUpb_MapContainer),
409 0,
410 Py_TPFLAGS_DEFAULT,
411 PyUpb_ScalarMapContainer_Slots,
412};
413
414// -----------------------------------------------------------------------------
415// MessageMapContainer
416// -----------------------------------------------------------------------------
417
418static PyMethodDef PyUpb_MessageMapContainer_Methods[] = {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700419 {"clear", PyUpb_MapContainer_Clear, METH_NOARGS,
420 "Removes all elements from the map."},
Jie Luo81da6b92024-11-12 09:41:02 -0800421 {"setdefault", (PyCFunction)PyUpb_MessageMapContainer_Setdefault,
422 METH_VARARGS, "setdefault is disallowed in MessageMap."},
Adam Cozzette501ecec2023-09-26 14:36:20 -0700423 {"get", (PyCFunction)PyUpb_MapContainer_Get, METH_VARARGS | METH_KEYWORDS,
424 "Gets the value for the given key if present, or otherwise a default"},
425 {"get_or_create", PyUpb_MapContainer_Subscript, METH_O,
426 "Alias for getitem, useful to make explicit that the map is mutated."},
427 {"GetEntryClass", PyUpb_MapContainer_GetEntryClass, METH_NOARGS,
428 "Return the class used to build Entries of (key, value) pairs."},
429 {"MergeFrom", PyUpb_MapContainer_MergeFrom, METH_O,
430 "Merges a map into the current map."},
431 /*
432 { "__deepcopy__", (PyCFunction)DeepCopy, METH_VARARGS,
433 "Makes a deep copy of the class." },
434 { "__reduce__", (PyCFunction)Reduce, METH_NOARGS,
435 "Outputs picklable representation of the repeated field." },
436 */
437 {NULL, NULL},
438};
439
440static PyType_Slot PyUpb_MessageMapContainer_Slots[] = {
441 {Py_tp_dealloc, PyUpb_MapContainer_Dealloc},
442 {Py_mp_length, PyUpb_MapContainer_Length},
443 {Py_mp_subscript, PyUpb_MapContainer_Subscript},
444 {Py_mp_ass_subscript, PyUpb_MapContainer_AssignSubscript},
Sandy Zhang5b329362024-03-05 17:54:43 -0800445 {Py_sq_contains, PyUpb_MapContainer_Contains},
Adam Cozzette501ecec2023-09-26 14:36:20 -0700446 {Py_tp_methods, PyUpb_MessageMapContainer_Methods},
447 {Py_tp_iter, PyUpb_MapIterator_New},
448 {Py_tp_repr, PyUpb_MapContainer_Repr},
449 {0, NULL}};
450
451static PyType_Spec PyUpb_MessageMapContainer_Spec = {
452 PYUPB_MODULE_NAME ".MessageMapContainer", sizeof(PyUpb_MapContainer), 0,
453 Py_TPFLAGS_DEFAULT, PyUpb_MessageMapContainer_Slots};
454
455// -----------------------------------------------------------------------------
456// MapIterator
457// -----------------------------------------------------------------------------
458
459typedef struct {
460 PyObject_HEAD;
461 PyUpb_MapContainer* map; // We own a reference.
462 size_t iter;
463 int version;
464} PyUpb_MapIterator;
465
466static PyObject* PyUpb_MapIterator_New(PyUpb_MapContainer* map) {
467 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
468 PyUpb_MapIterator* iter =
469 (void*)PyType_GenericAlloc(state->map_iterator_type, 0);
470 iter->map = map;
471 iter->iter = kUpb_Map_Begin;
472 iter->version = map->version;
473 Py_INCREF(map);
474 return &iter->ob_base;
475}
476
477static void PyUpb_MapIterator_Dealloc(void* _self) {
478 PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
479 Py_DECREF(&self->map->ob_base);
480 PyUpb_Dealloc(_self);
481}
482
Eric Salo33cb42e2023-12-19 17:48:33 -0800483static PyObject* PyUpb_MapIterator_IterNext(PyObject* _self) {
Adam Cozzette501ecec2023-09-26 14:36:20 -0700484 PyUpb_MapIterator* self = (PyUpb_MapIterator*)_self;
485 if (self->version != self->map->version) {
486 return PyErr_Format(PyExc_RuntimeError, "Map modified during iteration.");
487 }
488 upb_Map* map = PyUpb_MapContainer_GetIfReified(self->map);
489 if (!map) return NULL;
490 upb_MessageValue key, val;
491 if (!upb_Map_Next(map, &key, &val, &self->iter)) return NULL;
492 const upb_FieldDef* f = PyUpb_MapContainer_GetField(self->map);
493 const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
494 const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
495 return PyUpb_UpbToPy(key, key_f, self->map->arena);
496}
497
498static PyType_Slot PyUpb_MapIterator_Slots[] = {
499 {Py_tp_dealloc, PyUpb_MapIterator_Dealloc},
500 {Py_tp_iter, PyObject_SelfIter},
501 {Py_tp_iternext, PyUpb_MapIterator_IterNext},
502 {0, NULL}};
503
504static PyType_Spec PyUpb_MapIterator_Spec = {
505 PYUPB_MODULE_NAME ".MapIterator", sizeof(PyUpb_MapIterator), 0,
506 Py_TPFLAGS_DEFAULT, PyUpb_MapIterator_Slots};
507
508// -----------------------------------------------------------------------------
509// Top Level
510// -----------------------------------------------------------------------------
511
512static PyObject* GetMutableMappingBase(void) {
513 PyObject* collections = NULL;
514 PyObject* mapping = NULL;
Sandy Zhang5b329362024-03-05 17:54:43 -0800515 PyObject* base = NULL;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700516 if ((collections = PyImport_ImportModule("collections.abc")) &&
517 (mapping = PyObject_GetAttrString(collections, "MutableMapping"))) {
Sandy Zhang5b329362024-03-05 17:54:43 -0800518 base = Py_BuildValue("O", mapping);
Adam Cozzette501ecec2023-09-26 14:36:20 -0700519 }
520 Py_XDECREF(collections);
521 Py_XDECREF(mapping);
Sandy Zhang5b329362024-03-05 17:54:43 -0800522 return base;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700523}
524
525bool PyUpb_Map_Init(PyObject* m) {
526 PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
Sandy Zhang5b329362024-03-05 17:54:43 -0800527 PyObject* base = GetMutableMappingBase();
528 if (!base) return false;
Adam Cozzette501ecec2023-09-26 14:36:20 -0700529
Jie Luo81da6b92024-11-12 09:41:02 -0800530 const char* methods[] = {"keys", "items", "values", "__eq__", "__ne__",
531 "pop", "popitem", "update", NULL};
Sandy Zhang5b329362024-03-05 17:54:43 -0800532
533 state->message_map_container_type = PyUpb_AddClassWithRegister(
534 m, &PyUpb_MessageMapContainer_Spec, base, methods);
535 if (!state->message_map_container_type) {
536 return false;
537 }
538 state->scalar_map_container_type = PyUpb_AddClassWithRegister(
539 m, &PyUpb_ScalarMapContainer_Spec, base, methods);
540 if (!state->scalar_map_container_type) {
541 return false;
542 }
Adam Cozzette501ecec2023-09-26 14:36:20 -0700543 state->map_iterator_type = PyUpb_AddClass(m, &PyUpb_MapIterator_Spec);
544
Sandy Zhang5b329362024-03-05 17:54:43 -0800545 Py_DECREF(base);
546 Py_DECREF(methods);
Adam Cozzette501ecec2023-09-26 14:36:20 -0700547
548 return state->message_map_container_type &&
549 state->scalar_map_container_type && state->map_iterator_type;
550}