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/descriptor_containers.h"
29
30 #include "python/descriptor.h"
31 #include "python/protobuf.h"
32 #include "upb/reflection/def.h"
33
34 // Implements __repr__ as str(dict(self)).
PyUpb_DescriptorMap_Repr(PyObject * _self)35 static PyObject* PyUpb_DescriptorMap_Repr(PyObject* _self) {
36 PyObject* dict = PyDict_New();
37 PyObject* ret = NULL;
38 if (!dict) goto err;
39 if (PyDict_Merge(dict, _self, 1) != 0) goto err;
40 ret = PyObject_Str(dict);
41
42 err:
43 Py_XDECREF(dict);
44 return ret;
45 }
46
47 // -----------------------------------------------------------------------------
48 // ByNameIterator
49 // -----------------------------------------------------------------------------
50
51 typedef struct {
52 PyObject_HEAD;
53 const PyUpb_ByNameMap_Funcs* funcs;
54 const void* parent; // upb_MessageDef*, upb_DefPool*, etc.
55 PyObject* parent_obj; // Python object that keeps parent alive, we own a ref.
56 int index; // Current iterator index.
57 } PyUpb_ByNameIterator;
58
PyUpb_ByNameIterator_Self(PyObject * obj)59 static PyUpb_ByNameIterator* PyUpb_ByNameIterator_Self(PyObject* obj) {
60 assert(Py_TYPE(obj) == PyUpb_ModuleState_Get()->by_name_iterator_type);
61 return (PyUpb_ByNameIterator*)obj;
62 }
63
PyUpb_ByNameIterator_Dealloc(PyObject * _self)64 static void PyUpb_ByNameIterator_Dealloc(PyObject* _self) {
65 PyUpb_ByNameIterator* self = PyUpb_ByNameIterator_Self(_self);
66 Py_DECREF(self->parent_obj);
67 PyUpb_Dealloc(self);
68 }
69
PyUpb_ByNameIterator_New(const PyUpb_ByNameMap_Funcs * funcs,const void * parent,PyObject * parent_obj)70 static PyObject* PyUpb_ByNameIterator_New(const PyUpb_ByNameMap_Funcs* funcs,
71 const void* parent,
72 PyObject* parent_obj) {
73 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
74 PyUpb_ByNameIterator* iter =
75 (void*)PyType_GenericAlloc(s->by_name_iterator_type, 0);
76 iter->funcs = funcs;
77 iter->parent = parent;
78 iter->parent_obj = parent_obj;
79 iter->index = 0;
80 Py_INCREF(iter->parent_obj);
81 return &iter->ob_base;
82 }
83
PyUpb_ByNameIterator_IterNext(PyObject * _self)84 static PyObject* PyUpb_ByNameIterator_IterNext(PyObject* _self) {
85 PyUpb_ByNameIterator* self = PyUpb_ByNameIterator_Self(_self);
86 int size = self->funcs->base.get_elem_count(self->parent);
87 if (self->index >= size) return NULL;
88 const void* elem = self->funcs->base.index(self->parent, self->index);
89 self->index++;
90 return PyUnicode_FromString(self->funcs->get_elem_name(elem));
91 }
92
93 static PyType_Slot PyUpb_ByNameIterator_Slots[] = {
94 {Py_tp_dealloc, PyUpb_ByNameIterator_Dealloc},
95 {Py_tp_iter, PyObject_SelfIter},
96 {Py_tp_iternext, PyUpb_ByNameIterator_IterNext},
97 {0, NULL}};
98
99 static PyType_Spec PyUpb_ByNameIterator_Spec = {
100 PYUPB_MODULE_NAME "._ByNameIterator", // tp_name
101 sizeof(PyUpb_ByNameIterator), // tp_basicsize
102 0, // tp_itemsize
103 Py_TPFLAGS_DEFAULT, // tp_flags
104 PyUpb_ByNameIterator_Slots,
105 };
106
107 // -----------------------------------------------------------------------------
108 // ByNumberIterator
109 // -----------------------------------------------------------------------------
110
111 typedef struct {
112 PyObject_HEAD;
113 const PyUpb_ByNumberMap_Funcs* funcs;
114 const void* parent; // upb_MessageDef*, upb_DefPool*, etc.
115 PyObject* parent_obj; // Python object that keeps parent alive, we own a ref.
116 int index; // Current iterator index.
117 } PyUpb_ByNumberIterator;
118
PyUpb_ByNumberIterator_Self(PyObject * obj)119 static PyUpb_ByNumberIterator* PyUpb_ByNumberIterator_Self(PyObject* obj) {
120 assert(Py_TYPE(obj) == PyUpb_ModuleState_Get()->by_number_iterator_type);
121 return (PyUpb_ByNumberIterator*)obj;
122 }
123
PyUpb_ByNumberIterator_Dealloc(PyObject * _self)124 static void PyUpb_ByNumberIterator_Dealloc(PyObject* _self) {
125 PyUpb_ByNumberIterator* self = PyUpb_ByNumberIterator_Self(_self);
126 Py_DECREF(self->parent_obj);
127 PyUpb_Dealloc(self);
128 }
129
PyUpb_ByNumberIterator_New(const PyUpb_ByNumberMap_Funcs * funcs,const void * parent,PyObject * parent_obj)130 static PyObject* PyUpb_ByNumberIterator_New(
131 const PyUpb_ByNumberMap_Funcs* funcs, const void* parent,
132 PyObject* parent_obj) {
133 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
134 PyUpb_ByNumberIterator* iter =
135 (void*)PyType_GenericAlloc(s->by_number_iterator_type, 0);
136 iter->funcs = funcs;
137 iter->parent = parent;
138 iter->parent_obj = parent_obj;
139 iter->index = 0;
140 Py_INCREF(iter->parent_obj);
141 return &iter->ob_base;
142 }
143
PyUpb_ByNumberIterator_IterNext(PyObject * _self)144 static PyObject* PyUpb_ByNumberIterator_IterNext(PyObject* _self) {
145 PyUpb_ByNumberIterator* self = PyUpb_ByNumberIterator_Self(_self);
146 int size = self->funcs->base.get_elem_count(self->parent);
147 if (self->index >= size) return NULL;
148 const void* elem = self->funcs->base.index(self->parent, self->index);
149 self->index++;
150 return PyLong_FromLong(self->funcs->get_elem_num(elem));
151 }
152
153 static PyType_Slot PyUpb_ByNumberIterator_Slots[] = {
154 {Py_tp_dealloc, PyUpb_ByNumberIterator_Dealloc},
155 {Py_tp_iter, PyObject_SelfIter},
156 {Py_tp_iternext, PyUpb_ByNumberIterator_IterNext},
157 {0, NULL}};
158
159 static PyType_Spec PyUpb_ByNumberIterator_Spec = {
160 PYUPB_MODULE_NAME "._ByNumberIterator", // tp_name
161 sizeof(PyUpb_ByNumberIterator), // tp_basicsize
162 0, // tp_itemsize
163 Py_TPFLAGS_DEFAULT, // tp_flags
164 PyUpb_ByNumberIterator_Slots,
165 };
166
167 // -----------------------------------------------------------------------------
168 // GenericSequence
169 // -----------------------------------------------------------------------------
170
171 typedef struct {
172 PyObject_HEAD;
173 const PyUpb_GenericSequence_Funcs* funcs;
174 const void* parent; // upb_MessageDef*, upb_DefPool*, etc.
175 PyObject* parent_obj; // Python object that keeps parent alive, we own a ref.
176 } PyUpb_GenericSequence;
177
PyUpb_GenericSequence_Self(PyObject * obj)178 PyUpb_GenericSequence* PyUpb_GenericSequence_Self(PyObject* obj) {
179 assert(Py_TYPE(obj) == PyUpb_ModuleState_Get()->generic_sequence_type);
180 return (PyUpb_GenericSequence*)obj;
181 }
182
PyUpb_GenericSequence_Dealloc(PyObject * _self)183 static void PyUpb_GenericSequence_Dealloc(PyObject* _self) {
184 PyUpb_GenericSequence* self = PyUpb_GenericSequence_Self(_self);
185 Py_CLEAR(self->parent_obj);
186 PyUpb_Dealloc(self);
187 }
188
PyUpb_GenericSequence_New(const PyUpb_GenericSequence_Funcs * funcs,const void * parent,PyObject * parent_obj)189 PyObject* PyUpb_GenericSequence_New(const PyUpb_GenericSequence_Funcs* funcs,
190 const void* parent, PyObject* parent_obj) {
191 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
192 PyUpb_GenericSequence* seq =
193 (PyUpb_GenericSequence*)PyType_GenericAlloc(s->generic_sequence_type, 0);
194 seq->funcs = funcs;
195 seq->parent = parent;
196 seq->parent_obj = parent_obj;
197 Py_INCREF(parent_obj);
198 return &seq->ob_base;
199 }
200
PyUpb_GenericSequence_Length(PyObject * _self)201 static Py_ssize_t PyUpb_GenericSequence_Length(PyObject* _self) {
202 PyUpb_GenericSequence* self = PyUpb_GenericSequence_Self(_self);
203 return self->funcs->get_elem_count(self->parent);
204 }
205
PyUpb_GenericSequence_GetItem(PyObject * _self,Py_ssize_t index)206 static PyObject* PyUpb_GenericSequence_GetItem(PyObject* _self,
207 Py_ssize_t index) {
208 PyUpb_GenericSequence* self = PyUpb_GenericSequence_Self(_self);
209 Py_ssize_t size = self->funcs->get_elem_count(self->parent);
210 if (index < 0 || index >= size) {
211 PyErr_Format(PyExc_IndexError, "list index (%zd) out of range", index);
212 return NULL;
213 }
214 const void* elem = self->funcs->index(self->parent, index);
215 return self->funcs->get_elem_wrapper(elem);
216 }
217
218 // A sequence container can only be equal to another sequence container, or (for
219 // backward compatibility) to a list containing the same items.
220 // Returns 1 if equal, 0 if unequal, -1 on error.
PyUpb_GenericSequence_IsEqual(PyUpb_GenericSequence * self,PyObject * other)221 static int PyUpb_GenericSequence_IsEqual(PyUpb_GenericSequence* self,
222 PyObject* other) {
223 // Check the identity of C++ pointers.
224 if (PyObject_TypeCheck(other, Py_TYPE(self))) {
225 PyUpb_GenericSequence* other_seq = (void*)other;
226 return self->parent == other_seq->parent && self->funcs == other_seq->funcs;
227 }
228
229 if (!PyList_Check(other)) return 0;
230
231 // return list(self) == other
232 // We can clamp `i` to int because GenericSequence uses int for size (this
233 // is useful when we do int iteration below).
234 int n = PyUpb_GenericSequence_Length((PyObject*)self);
235 if ((Py_ssize_t)n != PyList_Size(other)) {
236 return false;
237 }
238
239 PyObject* item1;
240 for (int i = 0; i < n; i++) {
241 item1 = PyUpb_GenericSequence_GetItem((PyObject*)self, i);
242 PyObject* item2 = PyList_GetItem(other, i);
243 if (!item1 || !item2) goto error;
244 int cmp = PyObject_RichCompareBool(item1, item2, Py_EQ);
245 Py_DECREF(item1);
246 if (cmp != 1) return cmp;
247 }
248 // All items were found and equal
249 return 1;
250
251 error:
252 Py_XDECREF(item1);
253 return -1;
254 }
255
PyUpb_GenericSequence_RichCompare(PyObject * _self,PyObject * other,int opid)256 static PyObject* PyUpb_GenericSequence_RichCompare(PyObject* _self,
257 PyObject* other, int opid) {
258 PyUpb_GenericSequence* self = PyUpb_GenericSequence_Self(_self);
259 if (opid != Py_EQ && opid != Py_NE) {
260 Py_RETURN_NOTIMPLEMENTED;
261 }
262 bool ret = PyUpb_GenericSequence_IsEqual(self, other);
263 if (opid == Py_NE) ret = !ret;
264 return PyBool_FromLong(ret);
265 }
266
267 // Linear search. Could optimize this in some cases (defs that have index),
268 // but not all (FileDescriptor.dependencies).
PyUpb_GenericSequence_Find(PyObject * _self,PyObject * item)269 static int PyUpb_GenericSequence_Find(PyObject* _self, PyObject* item) {
270 PyUpb_GenericSequence* self = PyUpb_GenericSequence_Self(_self);
271 const void* item_ptr = PyUpb_AnyDescriptor_GetDef(item);
272 int count = self->funcs->get_elem_count(self->parent);
273 for (int i = 0; i < count; i++) {
274 if (self->funcs->index(self->parent, i) == item_ptr) {
275 return i;
276 }
277 }
278 return -1;
279 }
280
PyUpb_GenericSequence_Index(PyObject * self,PyObject * item)281 static PyObject* PyUpb_GenericSequence_Index(PyObject* self, PyObject* item) {
282 int position = PyUpb_GenericSequence_Find(self, item);
283 if (position < 0) {
284 PyErr_SetNone(PyExc_ValueError);
285 return NULL;
286 } else {
287 return PyLong_FromLong(position);
288 }
289 }
290
PyUpb_GenericSequence_Count(PyObject * _self,PyObject * item)291 static PyObject* PyUpb_GenericSequence_Count(PyObject* _self, PyObject* item) {
292 PyUpb_GenericSequence* self = PyUpb_GenericSequence_Self(_self);
293 const void* item_ptr = PyUpb_AnyDescriptor_GetDef(item);
294 int n = self->funcs->get_elem_count(self->parent);
295 int count = 0;
296 for (int i = 0; i < n; i++) {
297 if (self->funcs->index(self->parent, i) == item_ptr) {
298 count++;
299 }
300 }
301 return PyLong_FromLong(count);
302 }
303
PyUpb_GenericSequence_Append(PyObject * self,PyObject * args)304 static PyObject* PyUpb_GenericSequence_Append(PyObject* self, PyObject* args) {
305 PyErr_Format(PyExc_TypeError, "'%R' is not a mutable sequence", self);
306 return NULL;
307 }
308
309 static PyMethodDef PyUpb_GenericSequence_Methods[] = {
310 {"index", PyUpb_GenericSequence_Index, METH_O},
311 {"count", PyUpb_GenericSequence_Count, METH_O},
312 {"append", PyUpb_GenericSequence_Append, METH_O},
313 // This was implemented for Python/C++ but so far has not been required.
314 //{ "__reversed__", (PyCFunction)Reversed, METH_NOARGS, },
315 {NULL}};
316
317 static PyType_Slot PyUpb_GenericSequence_Slots[] = {
318 {Py_tp_dealloc, &PyUpb_GenericSequence_Dealloc},
319 {Py_tp_methods, &PyUpb_GenericSequence_Methods},
320 {Py_sq_length, PyUpb_GenericSequence_Length},
321 {Py_sq_item, PyUpb_GenericSequence_GetItem},
322 {Py_tp_richcompare, &PyUpb_GenericSequence_RichCompare},
323 // These were implemented for Python/C++ but so far have not been required.
324 // {Py_tp_repr, &PyUpb_GenericSequence_Repr},
325 // {Py_sq_contains, PyUpb_GenericSequence_Contains},
326 // {Py_mp_subscript, PyUpb_GenericSequence_Subscript},
327 // {Py_mp_ass_subscript, PyUpb_GenericSequence_AssignSubscript},
328 {0, NULL},
329 };
330
331 static PyType_Spec PyUpb_GenericSequence_Spec = {
332 PYUPB_MODULE_NAME "._GenericSequence", // tp_name
333 sizeof(PyUpb_GenericSequence), // tp_basicsize
334 0, // tp_itemsize
335 Py_TPFLAGS_DEFAULT, // tp_flags
336 PyUpb_GenericSequence_Slots,
337 };
338
339 // -----------------------------------------------------------------------------
340 // ByNameMap
341 // -----------------------------------------------------------------------------
342
343 typedef struct {
344 PyObject_HEAD;
345 const PyUpb_ByNameMap_Funcs* funcs;
346 const void* parent; // upb_MessageDef*, upb_DefPool*, etc.
347 PyObject* parent_obj; // Python object that keeps parent alive, we own a ref.
348 } PyUpb_ByNameMap;
349
PyUpb_ByNameMap_Self(PyObject * obj)350 PyUpb_ByNameMap* PyUpb_ByNameMap_Self(PyObject* obj) {
351 assert(Py_TYPE(obj) == PyUpb_ModuleState_Get()->by_name_map_type);
352 return (PyUpb_ByNameMap*)obj;
353 }
354
PyUpb_ByNameMap_Dealloc(PyObject * _self)355 static void PyUpb_ByNameMap_Dealloc(PyObject* _self) {
356 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
357 Py_DECREF(self->parent_obj);
358 PyUpb_Dealloc(self);
359 }
360
PyUpb_ByNameMap_New(const PyUpb_ByNameMap_Funcs * funcs,const void * parent,PyObject * parent_obj)361 PyObject* PyUpb_ByNameMap_New(const PyUpb_ByNameMap_Funcs* funcs,
362 const void* parent, PyObject* parent_obj) {
363 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
364 PyUpb_ByNameMap* map = (void*)PyType_GenericAlloc(s->by_name_map_type, 0);
365 map->funcs = funcs;
366 map->parent = parent;
367 map->parent_obj = parent_obj;
368 Py_INCREF(parent_obj);
369 return &map->ob_base;
370 }
371
PyUpb_ByNameMap_Length(PyObject * _self)372 static Py_ssize_t PyUpb_ByNameMap_Length(PyObject* _self) {
373 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
374 return self->funcs->base.get_elem_count(self->parent);
375 }
376
PyUpb_ByNameMap_Subscript(PyObject * _self,PyObject * key)377 static PyObject* PyUpb_ByNameMap_Subscript(PyObject* _self, PyObject* key) {
378 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
379 const char* name = PyUpb_GetStrData(key);
380 const void* elem = name ? self->funcs->lookup(self->parent, name) : NULL;
381
382 if (!name && PyObject_Hash(key) == -1) return NULL;
383
384 if (elem) {
385 return self->funcs->base.get_elem_wrapper(elem);
386 } else {
387 PyErr_SetObject(PyExc_KeyError, key);
388 return NULL;
389 }
390 }
391
PyUpb_ByNameMap_AssignSubscript(PyObject * self,PyObject * key,PyObject * value)392 static int PyUpb_ByNameMap_AssignSubscript(PyObject* self, PyObject* key,
393 PyObject* value) {
394 PyErr_Format(PyExc_TypeError, PYUPB_MODULE_NAME
395 ".ByNameMap' object does not support item assignment");
396 return -1;
397 }
398
PyUpb_ByNameMap_Contains(PyObject * _self,PyObject * key)399 static int PyUpb_ByNameMap_Contains(PyObject* _self, PyObject* key) {
400 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
401 const char* name = PyUpb_GetStrData(key);
402 const void* elem = name ? self->funcs->lookup(self->parent, name) : NULL;
403 if (!name && PyObject_Hash(key) == -1) return -1;
404 return elem ? 1 : 0;
405 }
406
PyUpb_ByNameMap_Get(PyObject * _self,PyObject * args)407 static PyObject* PyUpb_ByNameMap_Get(PyObject* _self, PyObject* args) {
408 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
409 PyObject* key;
410 PyObject* default_value = Py_None;
411 if (!PyArg_UnpackTuple(args, "get", 1, 2, &key, &default_value)) {
412 return NULL;
413 }
414
415 const char* name = PyUpb_GetStrData(key);
416 const void* elem = name ? self->funcs->lookup(self->parent, name) : NULL;
417
418 if (!name && PyObject_Hash(key) == -1) return NULL;
419
420 if (elem) {
421 return self->funcs->base.get_elem_wrapper(elem);
422 } else {
423 Py_INCREF(default_value);
424 return default_value;
425 }
426 }
427
PyUpb_ByNameMap_GetIter(PyObject * _self)428 static PyObject* PyUpb_ByNameMap_GetIter(PyObject* _self) {
429 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
430 return PyUpb_ByNameIterator_New(self->funcs, self->parent, self->parent_obj);
431 }
432
PyUpb_ByNameMap_Keys(PyObject * _self,PyObject * args)433 static PyObject* PyUpb_ByNameMap_Keys(PyObject* _self, PyObject* args) {
434 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
435 int n = self->funcs->base.get_elem_count(self->parent);
436 PyObject* ret = PyList_New(n);
437 if (!ret) return NULL;
438 for (int i = 0; i < n; i++) {
439 const void* elem = self->funcs->base.index(self->parent, i);
440 PyObject* key = PyUnicode_FromString(self->funcs->get_elem_name(elem));
441 if (!key) goto error;
442 PyList_SetItem(ret, i, key);
443 }
444 return ret;
445
446 error:
447 Py_XDECREF(ret);
448 return NULL;
449 }
450
PyUpb_ByNameMap_Values(PyObject * _self,PyObject * args)451 static PyObject* PyUpb_ByNameMap_Values(PyObject* _self, PyObject* args) {
452 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
453 int n = self->funcs->base.get_elem_count(self->parent);
454 PyObject* ret = PyList_New(n);
455 if (!ret) return NULL;
456 for (int i = 0; i < n; i++) {
457 const void* elem = self->funcs->base.index(self->parent, i);
458 PyObject* py_elem = self->funcs->base.get_elem_wrapper(elem);
459 if (!py_elem) goto error;
460 PyList_SetItem(ret, i, py_elem);
461 }
462 return ret;
463
464 error:
465 Py_XDECREF(ret);
466 return NULL;
467 }
468
PyUpb_ByNameMap_Items(PyObject * _self,PyObject * args)469 static PyObject* PyUpb_ByNameMap_Items(PyObject* _self, PyObject* args) {
470 PyUpb_ByNameMap* self = (PyUpb_ByNameMap*)_self;
471 int n = self->funcs->base.get_elem_count(self->parent);
472 PyObject* ret = PyList_New(n);
473 PyObject* item;
474 PyObject* py_elem;
475 if (!ret) return NULL;
476 for (int i = 0; i < n; i++) {
477 const void* elem = self->funcs->base.index(self->parent, i);
478 item = PyTuple_New(2);
479 py_elem = self->funcs->base.get_elem_wrapper(elem);
480 if (!item || !py_elem) goto error;
481 PyTuple_SetItem(item, 0,
482 PyUnicode_FromString(self->funcs->get_elem_name(elem)));
483 PyTuple_SetItem(item, 1, py_elem);
484 PyList_SetItem(ret, i, item);
485 }
486 return ret;
487
488 error:
489 Py_XDECREF(py_elem);
490 Py_XDECREF(item);
491 Py_XDECREF(ret);
492 return NULL;
493 }
494
495 // A mapping container can only be equal to another mapping container, or (for
496 // backward compatibility) to a dict containing the same items.
497 // Returns 1 if equal, 0 if unequal, -1 on error.
PyUpb_ByNameMap_IsEqual(PyUpb_ByNameMap * self,PyObject * other)498 static int PyUpb_ByNameMap_IsEqual(PyUpb_ByNameMap* self, PyObject* other) {
499 // Check the identity of C++ pointers.
500 if (PyObject_TypeCheck(other, Py_TYPE(self))) {
501 PyUpb_ByNameMap* other_map = (void*)other;
502 return self->parent == other_map->parent && self->funcs == other_map->funcs;
503 }
504
505 if (!PyDict_Check(other)) return 0;
506
507 PyObject* self_dict = PyDict_New();
508 PyDict_Merge(self_dict, (PyObject*)self, 0);
509 int eq = PyObject_RichCompareBool(self_dict, other, Py_EQ);
510 Py_DECREF(self_dict);
511 return eq;
512 }
513
PyUpb_ByNameMap_RichCompare(PyObject * _self,PyObject * other,int opid)514 static PyObject* PyUpb_ByNameMap_RichCompare(PyObject* _self, PyObject* other,
515 int opid) {
516 PyUpb_ByNameMap* self = PyUpb_ByNameMap_Self(_self);
517 if (opid != Py_EQ && opid != Py_NE) {
518 Py_RETURN_NOTIMPLEMENTED;
519 }
520 bool ret = PyUpb_ByNameMap_IsEqual(self, other);
521 if (opid == Py_NE) ret = !ret;
522 return PyBool_FromLong(ret);
523 }
524
525 static PyMethodDef PyUpb_ByNameMap_Methods[] = {
526 {"get", (PyCFunction)&PyUpb_ByNameMap_Get, METH_VARARGS},
527 {"keys", PyUpb_ByNameMap_Keys, METH_NOARGS},
528 {"values", PyUpb_ByNameMap_Values, METH_NOARGS},
529 {"items", PyUpb_ByNameMap_Items, METH_NOARGS},
530 {NULL}};
531
532 static PyType_Slot PyUpb_ByNameMap_Slots[] = {
533 {Py_mp_ass_subscript, PyUpb_ByNameMap_AssignSubscript},
534 {Py_mp_length, PyUpb_ByNameMap_Length},
535 {Py_mp_subscript, PyUpb_ByNameMap_Subscript},
536 {Py_sq_contains, &PyUpb_ByNameMap_Contains},
537 {Py_tp_dealloc, &PyUpb_ByNameMap_Dealloc},
538 {Py_tp_iter, PyUpb_ByNameMap_GetIter},
539 {Py_tp_methods, &PyUpb_ByNameMap_Methods},
540 {Py_tp_repr, &PyUpb_DescriptorMap_Repr},
541 {Py_tp_richcompare, &PyUpb_ByNameMap_RichCompare},
542 {0, NULL},
543 };
544
545 static PyType_Spec PyUpb_ByNameMap_Spec = {
546 PYUPB_MODULE_NAME "._ByNameMap", // tp_name
547 sizeof(PyUpb_ByNameMap), // tp_basicsize
548 0, // tp_itemsize
549 Py_TPFLAGS_DEFAULT, // tp_flags
550 PyUpb_ByNameMap_Slots,
551 };
552
553 // -----------------------------------------------------------------------------
554 // ByNumberMap
555 // -----------------------------------------------------------------------------
556
557 typedef struct {
558 PyObject_HEAD;
559 const PyUpb_ByNumberMap_Funcs* funcs;
560 const void* parent; // upb_MessageDef*, upb_DefPool*, etc.
561 PyObject* parent_obj; // Python object that keeps parent alive, we own a ref.
562 } PyUpb_ByNumberMap;
563
PyUpb_ByNumberMap_Self(PyObject * obj)564 PyUpb_ByNumberMap* PyUpb_ByNumberMap_Self(PyObject* obj) {
565 assert(Py_TYPE(obj) == PyUpb_ModuleState_Get()->by_number_map_type);
566 return (PyUpb_ByNumberMap*)obj;
567 }
568
PyUpb_ByNumberMap_New(const PyUpb_ByNumberMap_Funcs * funcs,const void * parent,PyObject * parent_obj)569 PyObject* PyUpb_ByNumberMap_New(const PyUpb_ByNumberMap_Funcs* funcs,
570 const void* parent, PyObject* parent_obj) {
571 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
572 PyUpb_ByNumberMap* map = (void*)PyType_GenericAlloc(s->by_number_map_type, 0);
573 map->funcs = funcs;
574 map->parent = parent;
575 map->parent_obj = parent_obj;
576 Py_INCREF(parent_obj);
577 return &map->ob_base;
578 }
579
PyUpb_ByNumberMap_Dealloc(PyObject * _self)580 static void PyUpb_ByNumberMap_Dealloc(PyObject* _self) {
581 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
582 Py_DECREF(self->parent_obj);
583 PyUpb_Dealloc(self);
584 }
585
PyUpb_ByNumberMap_Length(PyObject * _self)586 static Py_ssize_t PyUpb_ByNumberMap_Length(PyObject* _self) {
587 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
588 return self->funcs->base.get_elem_count(self->parent);
589 }
590
PyUpb_ByNumberMap_LookupHelper(PyUpb_ByNumberMap * self,PyObject * key)591 static const void* PyUpb_ByNumberMap_LookupHelper(PyUpb_ByNumberMap* self,
592 PyObject* key) {
593 long num = PyLong_AsLong(key);
594 if (num == -1 && PyErr_Occurred()) {
595 PyErr_Clear();
596 // Ensure that the key is hashable (this will raise an error if not).
597 PyObject_Hash(key);
598 return NULL;
599 } else {
600 return self->funcs->lookup(self->parent, num);
601 }
602 }
603
PyUpb_ByNumberMap_Subscript(PyObject * _self,PyObject * key)604 static PyObject* PyUpb_ByNumberMap_Subscript(PyObject* _self, PyObject* key) {
605 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
606 const void* elem = PyUpb_ByNumberMap_LookupHelper(self, key);
607 if (elem) {
608 return self->funcs->base.get_elem_wrapper(elem);
609 } else {
610 if (!PyErr_Occurred()) {
611 PyErr_SetObject(PyExc_KeyError, key);
612 }
613 return NULL;
614 }
615 }
616
PyUpb_ByNumberMap_AssignSubscript(PyObject * self,PyObject * key,PyObject * value)617 static int PyUpb_ByNumberMap_AssignSubscript(PyObject* self, PyObject* key,
618 PyObject* value) {
619 PyErr_Format(PyExc_TypeError, PYUPB_MODULE_NAME
620 ".ByNumberMap' object does not support item assignment");
621 return -1;
622 }
623
PyUpb_ByNumberMap_Get(PyObject * _self,PyObject * args)624 static PyObject* PyUpb_ByNumberMap_Get(PyObject* _self, PyObject* args) {
625 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
626 PyObject* key;
627 PyObject* default_value = Py_None;
628 if (!PyArg_UnpackTuple(args, "get", 1, 2, &key, &default_value)) {
629 return NULL;
630 }
631
632 const void* elem = PyUpb_ByNumberMap_LookupHelper(self, key);
633 if (elem) {
634 return self->funcs->base.get_elem_wrapper(elem);
635 } else if (PyErr_Occurred()) {
636 return NULL;
637 } else {
638 return PyUpb_NewRef(default_value);
639 }
640 }
641
PyUpb_ByNumberMap_GetIter(PyObject * _self)642 static PyObject* PyUpb_ByNumberMap_GetIter(PyObject* _self) {
643 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
644 return PyUpb_ByNumberIterator_New(self->funcs, self->parent,
645 self->parent_obj);
646 }
647
PyUpb_ByNumberMap_Keys(PyObject * _self,PyObject * args)648 static PyObject* PyUpb_ByNumberMap_Keys(PyObject* _self, PyObject* args) {
649 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
650 int n = self->funcs->base.get_elem_count(self->parent);
651 PyObject* ret = PyList_New(n);
652 if (!ret) return NULL;
653 for (int i = 0; i < n; i++) {
654 const void* elem = self->funcs->base.index(self->parent, i);
655 PyObject* key = PyLong_FromLong(self->funcs->get_elem_num(elem));
656 if (!key) goto error;
657 PyList_SetItem(ret, i, key);
658 }
659 return ret;
660
661 error:
662 Py_XDECREF(ret);
663 return NULL;
664 }
665
PyUpb_ByNumberMap_Values(PyObject * _self,PyObject * args)666 static PyObject* PyUpb_ByNumberMap_Values(PyObject* _self, PyObject* args) {
667 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
668 int n = self->funcs->base.get_elem_count(self->parent);
669 PyObject* ret = PyList_New(n);
670 if (!ret) return NULL;
671 for (int i = 0; i < n; i++) {
672 const void* elem = self->funcs->base.index(self->parent, i);
673 PyObject* py_elem = self->funcs->base.get_elem_wrapper(elem);
674 if (!py_elem) goto error;
675 PyList_SetItem(ret, i, py_elem);
676 }
677 return ret;
678
679 error:
680 Py_XDECREF(ret);
681 return NULL;
682 }
683
PyUpb_ByNumberMap_Items(PyObject * _self,PyObject * args)684 static PyObject* PyUpb_ByNumberMap_Items(PyObject* _self, PyObject* args) {
685 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
686 int n = self->funcs->base.get_elem_count(self->parent);
687 PyObject* ret = PyList_New(n);
688 PyObject* item;
689 PyObject* py_elem;
690 if (!ret) return NULL;
691 for (int i = 0; i < n; i++) {
692 const void* elem = self->funcs->base.index(self->parent, i);
693 int number = self->funcs->get_elem_num(elem);
694 item = PyTuple_New(2);
695 py_elem = self->funcs->base.get_elem_wrapper(elem);
696 if (!item || !py_elem) goto error;
697 PyTuple_SetItem(item, 0, PyLong_FromLong(number));
698 PyTuple_SetItem(item, 1, py_elem);
699 PyList_SetItem(ret, i, item);
700 }
701 return ret;
702
703 error:
704 Py_XDECREF(py_elem);
705 Py_XDECREF(item);
706 Py_XDECREF(ret);
707 return NULL;
708 }
709
PyUpb_ByNumberMap_Contains(PyObject * _self,PyObject * key)710 static int PyUpb_ByNumberMap_Contains(PyObject* _self, PyObject* key) {
711 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
712 const void* elem = PyUpb_ByNumberMap_LookupHelper(self, key);
713 if (elem) return 1;
714 if (PyErr_Occurred()) return -1;
715 return 0;
716 }
717
718 // A mapping container can only be equal to another mapping container, or (for
719 // backward compatibility) to a dict containing the same items.
720 // Returns 1 if equal, 0 if unequal, -1 on error.
PyUpb_ByNumberMap_IsEqual(PyUpb_ByNumberMap * self,PyObject * other)721 static int PyUpb_ByNumberMap_IsEqual(PyUpb_ByNumberMap* self, PyObject* other) {
722 // Check the identity of C++ pointers.
723 if (PyObject_TypeCheck(other, Py_TYPE(self))) {
724 PyUpb_ByNumberMap* other_map = (void*)other;
725 return self->parent == other_map->parent && self->funcs == other_map->funcs;
726 }
727
728 if (!PyDict_Check(other)) return 0;
729
730 PyObject* self_dict = PyDict_New();
731 PyDict_Merge(self_dict, (PyObject*)self, 0);
732 int eq = PyObject_RichCompareBool(self_dict, other, Py_EQ);
733 Py_DECREF(self_dict);
734 return eq;
735 }
736
PyUpb_ByNumberMap_RichCompare(PyObject * _self,PyObject * other,int opid)737 static PyObject* PyUpb_ByNumberMap_RichCompare(PyObject* _self, PyObject* other,
738 int opid) {
739 PyUpb_ByNumberMap* self = PyUpb_ByNumberMap_Self(_self);
740 if (opid != Py_EQ && opid != Py_NE) {
741 Py_RETURN_NOTIMPLEMENTED;
742 }
743 bool ret = PyUpb_ByNumberMap_IsEqual(self, other);
744 if (opid == Py_NE) ret = !ret;
745 return PyBool_FromLong(ret);
746 }
747
748 static PyMethodDef PyUpb_ByNumberMap_Methods[] = {
749 {"get", (PyCFunction)&PyUpb_ByNumberMap_Get, METH_VARARGS},
750 {"keys", PyUpb_ByNumberMap_Keys, METH_NOARGS},
751 {"values", PyUpb_ByNumberMap_Values, METH_NOARGS},
752 {"items", PyUpb_ByNumberMap_Items, METH_NOARGS},
753 {NULL}};
754
755 static PyType_Slot PyUpb_ByNumberMap_Slots[] = {
756 {Py_mp_ass_subscript, PyUpb_ByNumberMap_AssignSubscript},
757 {Py_mp_length, PyUpb_ByNumberMap_Length},
758 {Py_mp_subscript, PyUpb_ByNumberMap_Subscript},
759 {Py_sq_contains, &PyUpb_ByNumberMap_Contains},
760 {Py_tp_dealloc, &PyUpb_ByNumberMap_Dealloc},
761 {Py_tp_iter, PyUpb_ByNumberMap_GetIter},
762 {Py_tp_methods, &PyUpb_ByNumberMap_Methods},
763 {Py_tp_repr, &PyUpb_DescriptorMap_Repr},
764 {Py_tp_richcompare, &PyUpb_ByNumberMap_RichCompare},
765 {0, NULL},
766 };
767
768 static PyType_Spec PyUpb_ByNumberMap_Spec = {
769 PYUPB_MODULE_NAME "._ByNumberMap", // tp_name
770 sizeof(PyUpb_ByNumberMap), // tp_basicsize
771 0, // tp_itemsize
772 Py_TPFLAGS_DEFAULT, // tp_flags
773 PyUpb_ByNumberMap_Slots,
774 };
775
776 // -----------------------------------------------------------------------------
777 // Top Level
778 // -----------------------------------------------------------------------------
779
PyUpb_InitDescriptorContainers(PyObject * m)780 bool PyUpb_InitDescriptorContainers(PyObject* m) {
781 PyUpb_ModuleState* s = PyUpb_ModuleState_GetFromModule(m);
782
783 s->by_name_map_type = PyUpb_AddClass(m, &PyUpb_ByNameMap_Spec);
784 s->by_number_map_type = PyUpb_AddClass(m, &PyUpb_ByNumberMap_Spec);
785 s->by_name_iterator_type = PyUpb_AddClass(m, &PyUpb_ByNameIterator_Spec);
786 s->by_number_iterator_type = PyUpb_AddClass(m, &PyUpb_ByNumberIterator_Spec);
787 s->generic_sequence_type = PyUpb_AddClass(m, &PyUpb_GenericSequence_Spec);
788
789 return s->by_name_map_type && s->by_number_map_type &&
790 s->by_name_iterator_type && s->by_number_iterator_type &&
791 s->generic_sequence_type;
792 }
793