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