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/repeated.h"
29 
30 #include "python/convert.h"
31 #include "python/message.h"
32 #include "python/protobuf.h"
33 
34 static PyObject* PyUpb_RepeatedCompositeContainer_Append(PyObject* _self,
35                                                          PyObject* value);
36 static PyObject* PyUpb_RepeatedScalarContainer_Append(PyObject* _self,
37                                                       PyObject* value);
38 
39 // For an expression like:
40 //    foo[index]
41 //
42 // Converts `index` to an effective i/count/step, for a repeated field
43 // field of size `size`.
IndexToRange(PyObject * index,Py_ssize_t size,Py_ssize_t * i,Py_ssize_t * count,Py_ssize_t * step)44 static bool IndexToRange(PyObject* index, Py_ssize_t size, Py_ssize_t* i,
45                          Py_ssize_t* count, Py_ssize_t* step) {
46   assert(i && count && step);
47   if (PySlice_Check(index)) {
48     Py_ssize_t start, stop;
49     if (PySlice_Unpack(index, &start, &stop, step) < 0) return false;
50     *count = PySlice_AdjustIndices(size, &start, &stop, *step);
51     *i = start;
52   } else {
53     *i = PyNumber_AsSsize_t(index, PyExc_IndexError);
54 
55     if (*i == -1 && PyErr_Occurred()) {
56       PyErr_SetString(PyExc_TypeError, "list indices must be integers");
57       return false;
58     }
59 
60     if (*i < 0) *i += size;
61     *step = 0;
62     *count = 1;
63 
64     if (*i < 0 || size <= *i) {
65       PyErr_Format(PyExc_IndexError, "list index out of range");
66       return false;
67     }
68   }
69   return true;
70 }
71 
72 // Wrapper for a repeated field.
73 typedef struct {
74   PyObject_HEAD;
75   PyObject* arena;
76   // The field descriptor (PyObject*).
77   // The low bit indicates whether the container is reified (see ptr below).
78   //   - low bit set: repeated field is a stub (no underlying data).
79   //   - low bit clear: repeated field is reified (points to upb_Array).
80   uintptr_t field;
81   union {
82     PyObject* parent;  // stub: owning pointer to parent message.
83     upb_Array* arr;    // reified: the data for this array.
84   } ptr;
85 } PyUpb_RepeatedContainer;
86 
PyUpb_RepeatedContainer_IsStub(PyUpb_RepeatedContainer * self)87 static bool PyUpb_RepeatedContainer_IsStub(PyUpb_RepeatedContainer* self) {
88   return self->field & 1;
89 }
90 
PyUpb_RepeatedContainer_GetFieldDescriptor(PyUpb_RepeatedContainer * self)91 static PyObject* PyUpb_RepeatedContainer_GetFieldDescriptor(
92     PyUpb_RepeatedContainer* self) {
93   return (PyObject*)(self->field & ~(uintptr_t)1);
94 }
95 
PyUpb_RepeatedContainer_GetField(PyUpb_RepeatedContainer * self)96 static const upb_FieldDef* PyUpb_RepeatedContainer_GetField(
97     PyUpb_RepeatedContainer* self) {
98   return PyUpb_FieldDescriptor_GetDef(
99       PyUpb_RepeatedContainer_GetFieldDescriptor(self));
100 }
101 
102 // If the repeated field is reified, returns it.  Otherwise, returns NULL.
103 // If NULL is returned, the object is empty and has no underlying data.
PyUpb_RepeatedContainer_GetIfReified(PyUpb_RepeatedContainer * self)104 static upb_Array* PyUpb_RepeatedContainer_GetIfReified(
105     PyUpb_RepeatedContainer* self) {
106   return PyUpb_RepeatedContainer_IsStub(self) ? NULL : self->ptr.arr;
107 }
108 
PyUpb_RepeatedContainer_Reify(PyObject * _self,upb_Array * arr)109 void PyUpb_RepeatedContainer_Reify(PyObject* _self, upb_Array* arr) {
110   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
111   assert(PyUpb_RepeatedContainer_IsStub(self));
112   if (!arr) {
113     const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
114     upb_Arena* arena = PyUpb_Arena_Get(self->arena);
115     arr = upb_Array_New(arena, upb_FieldDef_CType(f));
116   }
117   PyUpb_ObjCache_Add(arr, &self->ob_base);
118   Py_DECREF(self->ptr.parent);
119   self->ptr.arr = arr;  // Overwrites self->ptr.parent.
120   self->field &= ~(uintptr_t)1;
121   assert(!PyUpb_RepeatedContainer_IsStub(self));
122 }
123 
PyUpb_RepeatedContainer_EnsureReified(PyObject * _self)124 upb_Array* PyUpb_RepeatedContainer_EnsureReified(PyObject* _self) {
125   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
126   upb_Array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
127   if (arr) return arr;  // Already writable.
128 
129   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
130   upb_Arena* arena = PyUpb_Arena_Get(self->arena);
131   arr = upb_Array_New(arena, upb_FieldDef_CType(f));
132   PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f,
133                                   (upb_MessageValue){.array_val = arr});
134   PyUpb_RepeatedContainer_Reify((PyObject*)self, arr);
135   return arr;
136 }
137 
PyUpb_RepeatedContainer_Dealloc(PyObject * _self)138 static void PyUpb_RepeatedContainer_Dealloc(PyObject* _self) {
139   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
140   Py_DECREF(self->arena);
141   if (PyUpb_RepeatedContainer_IsStub(self)) {
142     PyUpb_Message_CacheDelete(self->ptr.parent,
143                               PyUpb_RepeatedContainer_GetField(self));
144     Py_DECREF(self->ptr.parent);
145   } else {
146     PyUpb_ObjCache_Delete(self->ptr.arr);
147   }
148   Py_DECREF(PyUpb_RepeatedContainer_GetFieldDescriptor(self));
149   PyUpb_Dealloc(self);
150 }
151 
PyUpb_RepeatedContainer_GetClass(const upb_FieldDef * f)152 static PyTypeObject* PyUpb_RepeatedContainer_GetClass(const upb_FieldDef* f) {
153   assert(upb_FieldDef_IsRepeated(f) && !upb_FieldDef_IsMap(f));
154   PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
155   return upb_FieldDef_IsSubMessage(f) ? state->repeated_composite_container_type
156                                       : state->repeated_scalar_container_type;
157 }
158 
PyUpb_RepeatedContainer_Length(PyObject * self)159 static Py_ssize_t PyUpb_RepeatedContainer_Length(PyObject* self) {
160   upb_Array* arr =
161       PyUpb_RepeatedContainer_GetIfReified((PyUpb_RepeatedContainer*)self);
162   return arr ? upb_Array_Size(arr) : 0;
163 }
164 
PyUpb_RepeatedContainer_NewStub(PyObject * parent,const upb_FieldDef * f,PyObject * arena)165 PyObject* PyUpb_RepeatedContainer_NewStub(PyObject* parent,
166                                           const upb_FieldDef* f,
167                                           PyObject* arena) {
168   // We only create stubs when the parent is reified, by convention.  However
169   // this is not an invariant: the parent could become reified at any time.
170   assert(PyUpb_Message_GetIfReified(parent) == NULL);
171   PyTypeObject* cls = PyUpb_RepeatedContainer_GetClass(f);
172   PyUpb_RepeatedContainer* repeated = (void*)PyType_GenericAlloc(cls, 0);
173   repeated->arena = arena;
174   repeated->field = (uintptr_t)PyUpb_FieldDescriptor_Get(f) | 1;
175   repeated->ptr.parent = parent;
176   Py_INCREF(arena);
177   Py_INCREF(parent);
178   return &repeated->ob_base;
179 }
180 
PyUpb_RepeatedContainer_GetOrCreateWrapper(upb_Array * arr,const upb_FieldDef * f,PyObject * arena)181 PyObject* PyUpb_RepeatedContainer_GetOrCreateWrapper(upb_Array* arr,
182                                                      const upb_FieldDef* f,
183                                                      PyObject* arena) {
184   PyObject* ret = PyUpb_ObjCache_Get(arr);
185   if (ret) return ret;
186 
187   PyTypeObject* cls = PyUpb_RepeatedContainer_GetClass(f);
188   PyUpb_RepeatedContainer* repeated = (void*)PyType_GenericAlloc(cls, 0);
189   repeated->arena = arena;
190   repeated->field = (uintptr_t)PyUpb_FieldDescriptor_Get(f);
191   repeated->ptr.arr = arr;
192   ret = &repeated->ob_base;
193   Py_INCREF(arena);
194   PyUpb_ObjCache_Add(arr, ret);
195   return ret;
196 }
197 
198 static PyObject* PyUpb_RepeatedContainer_MergeFrom(PyObject* _self,
199                                                    PyObject* args);
200 
PyUpb_RepeatedContainer_DeepCopy(PyObject * _self,PyObject * value)201 PyObject* PyUpb_RepeatedContainer_DeepCopy(PyObject* _self, PyObject* value) {
202   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
203   PyUpb_RepeatedContainer* clone =
204       (void*)PyType_GenericAlloc(Py_TYPE(_self), 0);
205   if (clone == NULL) return NULL;
206   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
207   clone->arena = PyUpb_Arena_New();
208   clone->field = (uintptr_t)PyUpb_FieldDescriptor_Get(f);
209   clone->ptr.arr =
210       upb_Array_New(PyUpb_Arena_Get(clone->arena), upb_FieldDef_CType(f));
211   PyUpb_ObjCache_Add(clone->ptr.arr, (PyObject*)clone);
212   PyObject* result = PyUpb_RepeatedContainer_MergeFrom((PyObject*)clone, _self);
213   if (!result) {
214     Py_DECREF(clone);
215     return NULL;
216   }
217   Py_DECREF(result);
218   return (PyObject*)clone;
219 }
220 
PyUpb_RepeatedContainer_Extend(PyObject * _self,PyObject * value)221 PyObject* PyUpb_RepeatedContainer_Extend(PyObject* _self, PyObject* value) {
222   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
223   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
224   size_t start_size = upb_Array_Size(arr);
225   PyObject* it = PyObject_GetIter(value);
226   if (!it) {
227     PyErr_SetString(PyExc_TypeError, "Value must be iterable");
228     return NULL;
229   }
230 
231   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
232   bool submsg = upb_FieldDef_IsSubMessage(f);
233   PyObject* e;
234 
235   while ((e = PyIter_Next(it))) {
236     PyObject* ret;
237     if (submsg) {
238       ret = PyUpb_RepeatedCompositeContainer_Append(_self, e);
239     } else {
240       ret = PyUpb_RepeatedScalarContainer_Append(_self, e);
241     }
242     Py_XDECREF(ret);
243     Py_DECREF(e);
244   }
245 
246   Py_DECREF(it);
247 
248   if (PyErr_Occurred()) {
249     upb_Array_Resize(arr, start_size, NULL);
250     return NULL;
251   }
252 
253   Py_RETURN_NONE;
254 }
255 
PyUpb_RepeatedContainer_Item(PyObject * _self,Py_ssize_t index)256 static PyObject* PyUpb_RepeatedContainer_Item(PyObject* _self,
257                                               Py_ssize_t index) {
258   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
259   upb_Array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
260   Py_ssize_t size = arr ? upb_Array_Size(arr) : 0;
261   if (index < 0 || index >= size) {
262     PyErr_Format(PyExc_IndexError, "list index (%zd) out of range", index);
263     return NULL;
264   }
265   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
266   return PyUpb_UpbToPy(upb_Array_Get(arr, index), f, self->arena);
267 }
268 
PyUpb_RepeatedContainer_ToList(PyObject * _self)269 PyObject* PyUpb_RepeatedContainer_ToList(PyObject* _self) {
270   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
271   upb_Array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
272   if (!arr) return PyList_New(0);
273 
274   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
275   size_t n = upb_Array_Size(arr);
276   PyObject* list = PyList_New(n);
277   for (size_t i = 0; i < n; i++) {
278     PyObject* val = PyUpb_UpbToPy(upb_Array_Get(arr, i), f, self->arena);
279     if (!val) {
280       Py_DECREF(list);
281       return NULL;
282     }
283     PyList_SetItem(list, i, val);
284   }
285   return list;
286 }
287 
PyUpb_RepeatedContainer_Repr(PyObject * _self)288 static PyObject* PyUpb_RepeatedContainer_Repr(PyObject* _self) {
289   PyObject* list = PyUpb_RepeatedContainer_ToList(_self);
290   if (!list) return NULL;
291   assert(!PyErr_Occurred());
292   PyObject* repr = PyObject_Repr(list);
293   Py_DECREF(list);
294   return repr;
295 }
296 
PyUpb_RepeatedContainer_RichCompare(PyObject * _self,PyObject * _other,int opid)297 static PyObject* PyUpb_RepeatedContainer_RichCompare(PyObject* _self,
298                                                      PyObject* _other,
299                                                      int opid) {
300   if (opid != Py_EQ && opid != Py_NE) {
301     Py_INCREF(Py_NotImplemented);
302     return Py_NotImplemented;
303   }
304   PyObject* list1 = PyUpb_RepeatedContainer_ToList(_self);
305   PyObject* list2 = _other;
306   PyObject* del = NULL;
307   if (PyObject_TypeCheck(_other, _self->ob_type)) {
308     del = list2 = PyUpb_RepeatedContainer_ToList(_other);
309   }
310   PyObject* ret = PyObject_RichCompare(list1, list2, opid);
311   Py_DECREF(list1);
312   Py_XDECREF(del);
313   return ret;
314 }
315 
PyUpb_RepeatedContainer_Subscript(PyObject * _self,PyObject * key)316 static PyObject* PyUpb_RepeatedContainer_Subscript(PyObject* _self,
317                                                    PyObject* key) {
318   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
319   upb_Array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
320   Py_ssize_t size = arr ? upb_Array_Size(arr) : 0;
321   Py_ssize_t idx, count, step;
322   if (!IndexToRange(key, size, &idx, &count, &step)) return NULL;
323   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
324   if (step == 0) {
325     return PyUpb_UpbToPy(upb_Array_Get(arr, idx), f, self->arena);
326   } else {
327     PyObject* list = PyList_New(count);
328     for (Py_ssize_t i = 0; i < count; i++, idx += step) {
329       upb_MessageValue msgval = upb_Array_Get(self->ptr.arr, idx);
330       PyObject* item = PyUpb_UpbToPy(msgval, f, self->arena);
331       if (!item) {
332         Py_DECREF(list);
333         return NULL;
334       }
335       PyList_SetItem(list, i, item);
336     }
337     return list;
338   }
339 }
340 
PyUpb_RepeatedContainer_SetSubscript(PyUpb_RepeatedContainer * self,upb_Array * arr,const upb_FieldDef * f,Py_ssize_t idx,Py_ssize_t count,Py_ssize_t step,PyObject * value)341 static int PyUpb_RepeatedContainer_SetSubscript(
342     PyUpb_RepeatedContainer* self, upb_Array* arr, const upb_FieldDef* f,
343     Py_ssize_t idx, Py_ssize_t count, Py_ssize_t step, PyObject* value) {
344   upb_Arena* arena = PyUpb_Arena_Get(self->arena);
345   if (upb_FieldDef_IsSubMessage(f)) {
346     PyErr_SetString(PyExc_TypeError, "does not support assignment");
347     return -1;
348   }
349 
350   if (step == 0) {
351     // Set single value.
352     upb_MessageValue msgval;
353     if (!PyUpb_PyToUpb(value, f, &msgval, arena)) return -1;
354     upb_Array_Set(arr, idx, msgval);
355     return 0;
356   }
357 
358   // Set range.
359   PyObject* seq =
360       PySequence_Fast(value, "must assign iterable to extended slice");
361   PyObject* item = NULL;
362   int ret = -1;
363   if (!seq) goto err;
364   Py_ssize_t seq_size = PySequence_Size(seq);
365   if (seq_size != count) {
366     if (step == 1) {
367       // We must shift the tail elements (either right or left).
368       size_t tail = upb_Array_Size(arr) - (idx + count);
369       upb_Array_Resize(arr, idx + seq_size + tail, arena);
370       upb_Array_Move(arr, idx + seq_size, idx + count, tail);
371       count = seq_size;
372     } else {
373       PyErr_Format(PyExc_ValueError,
374                    "attempt to assign sequence of  %zd to extended slice "
375                    "of size %zd",
376                    seq_size, count);
377       goto err;
378     }
379   }
380   for (Py_ssize_t i = 0; i < count; i++, idx += step) {
381     upb_MessageValue msgval;
382     item = PySequence_GetItem(seq, i);
383     if (!item) goto err;
384     // XXX: if this fails we can leave the list partially mutated.
385     if (!PyUpb_PyToUpb(item, f, &msgval, arena)) goto err;
386     Py_DECREF(item);
387     item = NULL;
388     upb_Array_Set(arr, idx, msgval);
389   }
390   ret = 0;
391 
392 err:
393   Py_XDECREF(seq);
394   Py_XDECREF(item);
395   return ret;
396 }
397 
PyUpb_RepeatedContainer_DeleteSubscript(upb_Array * arr,Py_ssize_t idx,Py_ssize_t count,Py_ssize_t step)398 static int PyUpb_RepeatedContainer_DeleteSubscript(upb_Array* arr,
399                                                    Py_ssize_t idx,
400                                                    Py_ssize_t count,
401                                                    Py_ssize_t step) {
402   // Normalize direction: deletion is order-independent.
403   Py_ssize_t start = idx;
404   if (step < 0) {
405     Py_ssize_t end = start + step * (count - 1);
406     start = end;
407     step = -step;
408   }
409 
410   size_t dst = start;
411   size_t src;
412   if (step > 1) {
413     // Move elements between steps:
414     //
415     //        src
416     //         |
417     // |------X---X---X---X------------------------------|
418     //        |
419     //       dst           <-------- tail -------------->
420     src = start + 1;
421     for (Py_ssize_t i = 1; i < count; i++, dst += step - 1, src += step) {
422       upb_Array_Move(arr, dst, src, step);
423     }
424   } else {
425     src = start + count;
426   }
427 
428   // Move tail.
429   size_t tail = upb_Array_Size(arr) - src;
430   size_t new_size = dst + tail;
431   assert(new_size == upb_Array_Size(arr) - count);
432   upb_Array_Move(arr, dst, src, tail);
433   upb_Array_Resize(arr, new_size, NULL);
434   return 0;
435 }
436 
PyUpb_RepeatedContainer_AssignSubscript(PyObject * _self,PyObject * key,PyObject * value)437 static int PyUpb_RepeatedContainer_AssignSubscript(PyObject* _self,
438                                                    PyObject* key,
439                                                    PyObject* value) {
440   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
441   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
442   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
443   Py_ssize_t size = arr ? upb_Array_Size(arr) : 0;
444   Py_ssize_t idx, count, step;
445   if (!IndexToRange(key, size, &idx, &count, &step)) return -1;
446   if (value) {
447     return PyUpb_RepeatedContainer_SetSubscript(self, arr, f, idx, count, step,
448                                                 value);
449   } else {
450     return PyUpb_RepeatedContainer_DeleteSubscript(arr, idx, count, step);
451   }
452 }
453 
PyUpb_RepeatedContainer_Pop(PyObject * _self,PyObject * args)454 static PyObject* PyUpb_RepeatedContainer_Pop(PyObject* _self, PyObject* args) {
455   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
456   Py_ssize_t index = -1;
457   if (!PyArg_ParseTuple(args, "|n", &index)) return NULL;
458   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
459   size_t size = upb_Array_Size(arr);
460   if (index < 0) index += size;
461   if (index >= size) index = size - 1;
462   PyObject* ret = PyUpb_RepeatedContainer_Item(_self, index);
463   if (!ret) return NULL;
464   upb_Array_Delete(self->ptr.arr, index, 1);
465   return ret;
466 }
467 
PyUpb_RepeatedContainer_Remove(PyObject * _self,PyObject * value)468 static PyObject* PyUpb_RepeatedContainer_Remove(PyObject* _self,
469                                                 PyObject* value) {
470   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
471   Py_ssize_t match_index = -1;
472   Py_ssize_t n = PyUpb_RepeatedContainer_Length(_self);
473   for (Py_ssize_t i = 0; i < n; ++i) {
474     PyObject* elem = PyUpb_RepeatedContainer_Item(_self, i);
475     if (!elem) return NULL;
476     int eq = PyObject_RichCompareBool(elem, value, Py_EQ);
477     Py_DECREF(elem);
478     if (eq) {
479       match_index = i;
480       break;
481     }
482   }
483   if (match_index == -1) {
484     PyErr_SetString(PyExc_ValueError, "remove(x): x not in container");
485     return NULL;
486   }
487   if (PyUpb_RepeatedContainer_DeleteSubscript(arr, match_index, 1, 1) < 0) {
488     return NULL;
489   }
490   Py_RETURN_NONE;
491 }
492 
493 // A helper function used only for Sort().
PyUpb_RepeatedContainer_Assign(PyObject * _self,PyObject * list)494 static bool PyUpb_RepeatedContainer_Assign(PyObject* _self, PyObject* list) {
495   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
496   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
497   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
498   Py_ssize_t size = PyList_Size(list);
499   bool submsg = upb_FieldDef_IsSubMessage(f);
500   upb_Arena* arena = PyUpb_Arena_Get(self->arena);
501   for (Py_ssize_t i = 0; i < size; ++i) {
502     PyObject* obj = PyList_GetItem(list, i);
503     upb_MessageValue msgval;
504     if (submsg) {
505       msgval.msg_val = PyUpb_Message_GetIfReified(obj);
506       assert(msgval.msg_val);
507     } else {
508       if (!PyUpb_PyToUpb(obj, f, &msgval, arena)) return false;
509     }
510     upb_Array_Set(arr, i, msgval);
511   }
512   return true;
513 }
514 
PyUpb_RepeatedContainer_Sort(PyObject * pself,PyObject * args,PyObject * kwds)515 static PyObject* PyUpb_RepeatedContainer_Sort(PyObject* pself, PyObject* args,
516                                               PyObject* kwds) {
517   // Support the old sort_function argument for backwards
518   // compatibility.
519   if (kwds != NULL) {
520     PyObject* sort_func = PyDict_GetItemString(kwds, "sort_function");
521     if (sort_func != NULL) {
522       // Must set before deleting as sort_func is a borrowed reference
523       // and kwds might be the only thing keeping it alive.
524       if (PyDict_SetItemString(kwds, "cmp", sort_func) == -1) return NULL;
525       if (PyDict_DelItemString(kwds, "sort_function") == -1) return NULL;
526     }
527   }
528 
529   PyObject* ret = NULL;
530   PyObject* full_slice = NULL;
531   PyObject* list = NULL;
532   PyObject* m = NULL;
533   PyObject* res = NULL;
534   if ((full_slice = PySlice_New(NULL, NULL, NULL)) &&
535       (list = PyUpb_RepeatedContainer_Subscript(pself, full_slice)) &&
536       (m = PyObject_GetAttrString(list, "sort")) &&
537       (res = PyObject_Call(m, args, kwds)) &&
538       PyUpb_RepeatedContainer_Assign(pself, list)) {
539     Py_INCREF(Py_None);
540     ret = Py_None;
541   }
542 
543   Py_XDECREF(full_slice);
544   Py_XDECREF(list);
545   Py_XDECREF(m);
546   Py_XDECREF(res);
547   return ret;
548 }
549 
PyUpb_RepeatedContainer_Reverse(PyObject * _self)550 static PyObject* PyUpb_RepeatedContainer_Reverse(PyObject* _self) {
551   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
552   size_t n = upb_Array_Size(arr);
553   size_t half = n / 2;  // Rounds down.
554   for (size_t i = 0; i < half; i++) {
555     size_t i2 = n - i - 1;
556     upb_MessageValue val1 = upb_Array_Get(arr, i);
557     upb_MessageValue val2 = upb_Array_Get(arr, i2);
558     upb_Array_Set(arr, i, val2);
559     upb_Array_Set(arr, i2, val1);
560   }
561   Py_RETURN_NONE;
562 }
563 
PyUpb_RepeatedContainer_MergeFrom(PyObject * _self,PyObject * args)564 static PyObject* PyUpb_RepeatedContainer_MergeFrom(PyObject* _self,
565                                                    PyObject* args) {
566   return PyUpb_RepeatedContainer_Extend(_self, args);
567 }
568 
569 // -----------------------------------------------------------------------------
570 // RepeatedCompositeContainer
571 // -----------------------------------------------------------------------------
572 
PyUpb_RepeatedCompositeContainer_AppendNew(PyObject * _self)573 static PyObject* PyUpb_RepeatedCompositeContainer_AppendNew(PyObject* _self) {
574   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
575   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
576   if (!arr) return NULL;
577   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
578   upb_Arena* arena = PyUpb_Arena_Get(self->arena);
579   const upb_MessageDef* m = upb_FieldDef_MessageSubDef(f);
580   const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
581   upb_Message* msg = upb_Message_New(layout, arena);
582   upb_MessageValue msgval = {.msg_val = msg};
583   upb_Array_Append(arr, msgval, arena);
584   return PyUpb_Message_Get(msg, m, self->arena);
585 }
586 
PyUpb_RepeatedCompositeContainer_Add(PyObject * _self,PyObject * args,PyObject * kwargs)587 PyObject* PyUpb_RepeatedCompositeContainer_Add(PyObject* _self, PyObject* args,
588                                                PyObject* kwargs) {
589   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
590   PyObject* py_msg = PyUpb_RepeatedCompositeContainer_AppendNew(_self);
591   if (!py_msg) return NULL;
592   if (PyUpb_Message_InitAttributes(py_msg, args, kwargs) < 0) {
593     Py_DECREF(py_msg);
594     upb_Array_Delete(self->ptr.arr, upb_Array_Size(self->ptr.arr) - 1, 1);
595     return NULL;
596   }
597   return py_msg;
598 }
599 
PyUpb_RepeatedCompositeContainer_Append(PyObject * _self,PyObject * value)600 static PyObject* PyUpb_RepeatedCompositeContainer_Append(PyObject* _self,
601                                                          PyObject* value) {
602   if (!PyUpb_Message_Verify(value)) return NULL;
603   PyObject* py_msg = PyUpb_RepeatedCompositeContainer_AppendNew(_self);
604   if (!py_msg) return NULL;
605   PyObject* none = PyUpb_Message_MergeFrom(py_msg, value);
606   if (!none) {
607     Py_DECREF(py_msg);
608     return NULL;
609   }
610   Py_DECREF(none);
611   return py_msg;
612 }
613 
PyUpb_RepeatedContainer_Insert(PyObject * _self,PyObject * args)614 static PyObject* PyUpb_RepeatedContainer_Insert(PyObject* _self,
615                                                 PyObject* args) {
616   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
617   Py_ssize_t index;
618   PyObject* value;
619   if (!PyArg_ParseTuple(args, "nO", &index, &value)) return NULL;
620   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
621   if (!arr) return NULL;
622 
623   // Normalize index.
624   Py_ssize_t size = upb_Array_Size(arr);
625   if (index < 0) index += size;
626   if (index < 0) index = 0;
627   if (index > size) index = size;
628 
629   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
630   upb_MessageValue msgval;
631   upb_Arena* arena = PyUpb_Arena_Get(self->arena);
632   if (upb_FieldDef_IsSubMessage(f)) {
633     // Create message.
634     const upb_MessageDef* m = upb_FieldDef_MessageSubDef(f);
635     const upb_MiniTable* layout = upb_MessageDef_MiniTable(m);
636     upb_Message* msg = upb_Message_New(layout, arena);
637     PyObject* py_msg = PyUpb_Message_Get(msg, m, self->arena);
638     PyObject* ret = PyUpb_Message_MergeFrom(py_msg, value);
639     Py_DECREF(py_msg);
640     if (!ret) return NULL;
641     Py_DECREF(ret);
642     msgval.msg_val = msg;
643   } else {
644     if (!PyUpb_PyToUpb(value, f, &msgval, arena)) return NULL;
645   }
646 
647   upb_Array_Insert(arr, index, 1, arena);
648   upb_Array_Set(arr, index, msgval);
649 
650   Py_RETURN_NONE;
651 }
652 
653 static PyMethodDef PyUpb_RepeatedCompositeContainer_Methods[] = {
654     {"__deepcopy__", PyUpb_RepeatedContainer_DeepCopy, METH_VARARGS,
655      "Makes a deep copy of the class."},
656     {"add", (PyCFunction)PyUpb_RepeatedCompositeContainer_Add,
657      METH_VARARGS | METH_KEYWORDS, "Adds an object to the repeated container."},
658     {"append", PyUpb_RepeatedCompositeContainer_Append, METH_O,
659      "Appends a message to the end of the repeated container."},
660     {"insert", PyUpb_RepeatedContainer_Insert, METH_VARARGS,
661      "Inserts a message before the specified index."},
662     {"extend", PyUpb_RepeatedContainer_Extend, METH_O,
663      "Adds objects to the repeated container."},
664     {"pop", PyUpb_RepeatedContainer_Pop, METH_VARARGS,
665      "Removes an object from the repeated container and returns it."},
666     {"remove", PyUpb_RepeatedContainer_Remove, METH_O,
667      "Removes an object from the repeated container."},
668     {"sort", (PyCFunction)PyUpb_RepeatedContainer_Sort,
669      METH_VARARGS | METH_KEYWORDS, "Sorts the repeated container."},
670     {"reverse", (PyCFunction)PyUpb_RepeatedContainer_Reverse, METH_NOARGS,
671      "Reverses elements order of the repeated container."},
672     {"MergeFrom", PyUpb_RepeatedContainer_MergeFrom, METH_O,
673      "Adds objects to the repeated container."},
674     {NULL, NULL}};
675 
676 static PyType_Slot PyUpb_RepeatedCompositeContainer_Slots[] = {
677     {Py_tp_dealloc, PyUpb_RepeatedContainer_Dealloc},
678     {Py_tp_methods, PyUpb_RepeatedCompositeContainer_Methods},
679     {Py_sq_length, PyUpb_RepeatedContainer_Length},
680     {Py_sq_item, PyUpb_RepeatedContainer_Item},
681     {Py_mp_length, PyUpb_RepeatedContainer_Length},
682     {Py_tp_repr, PyUpb_RepeatedContainer_Repr},
683     {Py_mp_subscript, PyUpb_RepeatedContainer_Subscript},
684     {Py_mp_ass_subscript, PyUpb_RepeatedContainer_AssignSubscript},
685     {Py_tp_new, PyUpb_Forbidden_New},
686     {Py_tp_richcompare, PyUpb_RepeatedContainer_RichCompare},
687     {Py_tp_hash, PyObject_HashNotImplemented},
688     {0, NULL}};
689 
690 static PyType_Spec PyUpb_RepeatedCompositeContainer_Spec = {
691     PYUPB_MODULE_NAME ".RepeatedCompositeContainer",
692     sizeof(PyUpb_RepeatedContainer),
693     0,  // tp_itemsize
694     Py_TPFLAGS_DEFAULT,
695     PyUpb_RepeatedCompositeContainer_Slots,
696 };
697 
698 // -----------------------------------------------------------------------------
699 // RepeatedScalarContainer
700 // -----------------------------------------------------------------------------
701 
PyUpb_RepeatedScalarContainer_Append(PyObject * _self,PyObject * value)702 static PyObject* PyUpb_RepeatedScalarContainer_Append(PyObject* _self,
703                                                       PyObject* value) {
704   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
705   upb_Array* arr = PyUpb_RepeatedContainer_EnsureReified(_self);
706   upb_Arena* arena = PyUpb_Arena_Get(self->arena);
707   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
708   upb_MessageValue msgval;
709   if (!PyUpb_PyToUpb(value, f, &msgval, arena)) {
710     return NULL;
711   }
712   upb_Array_Append(arr, msgval, arena);
713   Py_RETURN_NONE;
714 }
715 
PyUpb_RepeatedScalarContainer_AssignItem(PyObject * _self,Py_ssize_t index,PyObject * item)716 static int PyUpb_RepeatedScalarContainer_AssignItem(PyObject* _self,
717                                                     Py_ssize_t index,
718                                                     PyObject* item) {
719   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
720   upb_Array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
721   Py_ssize_t size = arr ? upb_Array_Size(arr) : 0;
722   if (index < 0 || index >= size) {
723     PyErr_Format(PyExc_IndexError, "list index (%zd) out of range", index);
724     return -1;
725   }
726   const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
727   upb_MessageValue msgval;
728   upb_Arena* arena = PyUpb_Arena_Get(self->arena);
729   if (!PyUpb_PyToUpb(item, f, &msgval, arena)) {
730     return -1;
731   }
732   upb_Array_Set(self->ptr.arr, index, msgval);
733   return 0;
734 }
735 
PyUpb_RepeatedScalarContainer_Reduce(PyObject * unused_self,PyObject * unused_other)736 static PyObject* PyUpb_RepeatedScalarContainer_Reduce(PyObject* unused_self,
737                                                       PyObject* unused_other) {
738   PyObject* pickle_module = PyImport_ImportModule("pickle");
739   if (!pickle_module) return NULL;
740   PyObject* pickle_error = PyObject_GetAttrString(pickle_module, "PickleError");
741   Py_DECREF(pickle_module);
742   if (!pickle_error) return NULL;
743   PyErr_Format(pickle_error,
744                "can't pickle repeated message fields, convert to list first");
745   Py_DECREF(pickle_error);
746   return NULL;
747 }
748 
749 static PyMethodDef PyUpb_RepeatedScalarContainer_Methods[] = {
750     {"__deepcopy__", PyUpb_RepeatedContainer_DeepCopy, METH_VARARGS,
751      "Makes a deep copy of the class."},
752     {"__reduce__", PyUpb_RepeatedScalarContainer_Reduce, METH_NOARGS,
753      "Outputs picklable representation of the repeated field."},
754     {"append", PyUpb_RepeatedScalarContainer_Append, METH_O,
755      "Appends an object to the repeated container."},
756     {"extend", PyUpb_RepeatedContainer_Extend, METH_O,
757      "Appends objects to the repeated container."},
758     {"insert", PyUpb_RepeatedContainer_Insert, METH_VARARGS,
759      "Inserts an object at the specified position in the container."},
760     {"pop", PyUpb_RepeatedContainer_Pop, METH_VARARGS,
761      "Removes an object from the repeated container and returns it."},
762     {"remove", PyUpb_RepeatedContainer_Remove, METH_O,
763      "Removes an object from the repeated container."},
764     {"sort", (PyCFunction)PyUpb_RepeatedContainer_Sort,
765      METH_VARARGS | METH_KEYWORDS, "Sorts the repeated container."},
766     {"reverse", (PyCFunction)PyUpb_RepeatedContainer_Reverse, METH_NOARGS,
767      "Reverses elements order of the repeated container."},
768     {"MergeFrom", PyUpb_RepeatedContainer_MergeFrom, METH_O,
769      "Merges a repeated container into the current container."},
770     {NULL, NULL}};
771 
772 static PyType_Slot PyUpb_RepeatedScalarContainer_Slots[] = {
773     {Py_tp_dealloc, PyUpb_RepeatedContainer_Dealloc},
774     {Py_tp_methods, PyUpb_RepeatedScalarContainer_Methods},
775     {Py_tp_new, PyUpb_Forbidden_New},
776     {Py_tp_repr, PyUpb_RepeatedContainer_Repr},
777     {Py_sq_length, PyUpb_RepeatedContainer_Length},
778     {Py_sq_item, PyUpb_RepeatedContainer_Item},
779     {Py_sq_ass_item, PyUpb_RepeatedScalarContainer_AssignItem},
780     {Py_mp_length, PyUpb_RepeatedContainer_Length},
781     {Py_mp_subscript, PyUpb_RepeatedContainer_Subscript},
782     {Py_mp_ass_subscript, PyUpb_RepeatedContainer_AssignSubscript},
783     {Py_tp_richcompare, PyUpb_RepeatedContainer_RichCompare},
784     {Py_tp_hash, PyObject_HashNotImplemented},
785     {0, NULL}};
786 
787 static PyType_Spec PyUpb_RepeatedScalarContainer_Spec = {
788     PYUPB_MODULE_NAME ".RepeatedScalarContainer",
789     sizeof(PyUpb_RepeatedContainer),
790     0,  // tp_itemsize
791     Py_TPFLAGS_DEFAULT,
792     PyUpb_RepeatedScalarContainer_Slots,
793 };
794 
795 // -----------------------------------------------------------------------------
796 // Top Level
797 // -----------------------------------------------------------------------------
798 
PyUpb_Repeated_RegisterAsSequence(PyUpb_ModuleState * state)799 static bool PyUpb_Repeated_RegisterAsSequence(PyUpb_ModuleState* state) {
800   PyObject* collections = NULL;
801   PyObject* seq = NULL;
802   PyObject* ret1 = NULL;
803   PyObject* ret2 = NULL;
804   PyTypeObject* type1 = state->repeated_scalar_container_type;
805   PyTypeObject* type2 = state->repeated_composite_container_type;
806 
807   bool ok = (collections = PyImport_ImportModule("collections.abc")) &&
808             (seq = PyObject_GetAttrString(collections, "MutableSequence")) &&
809             (ret1 = PyObject_CallMethod(seq, "register", "O", type1)) &&
810             (ret2 = PyObject_CallMethod(seq, "register", "O", type2));
811 
812   Py_XDECREF(collections);
813   Py_XDECREF(seq);
814   Py_XDECREF(ret1);
815   Py_XDECREF(ret2);
816   return ok;
817 }
818 
PyUpb_Repeated_Init(PyObject * m)819 bool PyUpb_Repeated_Init(PyObject* m) {
820   PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
821 
822   state->repeated_composite_container_type =
823       PyUpb_AddClass(m, &PyUpb_RepeatedCompositeContainer_Spec);
824   state->repeated_scalar_container_type =
825       PyUpb_AddClass(m, &PyUpb_RepeatedScalarContainer_Spec);
826 
827   return state->repeated_composite_container_type &&
828          state->repeated_scalar_container_type &&
829          PyUpb_Repeated_RegisterAsSequence(state);
830 }
831