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