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/extension_dict.h"
29
30 #include "python/message.h"
31 #include "python/protobuf.h"
32 #include "upb/reflection/def.h"
33
34 // -----------------------------------------------------------------------------
35 // ExtensionDict
36 // -----------------------------------------------------------------------------
37
38 typedef struct {
39 PyObject_HEAD;
40 PyObject* msg; // Owning ref to our parent pessage.
41 } PyUpb_ExtensionDict;
42
PyUpb_ExtensionDict_New(PyObject * msg)43 PyObject* PyUpb_ExtensionDict_New(PyObject* msg) {
44 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
45 PyUpb_ExtensionDict* ext_dict =
46 (void*)PyType_GenericAlloc(state->extension_dict_type, 0);
47 ext_dict->msg = msg;
48 Py_INCREF(ext_dict->msg);
49 return &ext_dict->ob_base;
50 }
51
PyUpb_ExtensionDict_FindExtensionByName(PyObject * _self,PyObject * key)52 static PyObject* PyUpb_ExtensionDict_FindExtensionByName(PyObject* _self,
53 PyObject* key) {
54 PyUpb_ExtensionDict* self = (PyUpb_ExtensionDict*)_self;
55 const char* name = PyUpb_GetStrData(key);
56 const upb_MessageDef* m = PyUpb_Message_GetMsgdef(self->msg);
57 const upb_FileDef* file = upb_MessageDef_File(m);
58 const upb_DefPool* symtab = upb_FileDef_Pool(file);
59 const upb_FieldDef* ext = upb_DefPool_FindExtensionByName(symtab, name);
60 if (ext) {
61 return PyUpb_FieldDescriptor_Get(ext);
62 } else {
63 Py_RETURN_NONE;
64 }
65 }
66
PyUpb_ExtensionDict_FindExtensionByNumber(PyObject * _self,PyObject * arg)67 static PyObject* PyUpb_ExtensionDict_FindExtensionByNumber(PyObject* _self,
68 PyObject* arg) {
69 PyUpb_ExtensionDict* self = (PyUpb_ExtensionDict*)_self;
70 const upb_MessageDef* m = PyUpb_Message_GetMsgdef(self->msg);
71 const upb_MiniTable* l = upb_MessageDef_MiniTable(m);
72 const upb_FileDef* file = upb_MessageDef_File(m);
73 const upb_DefPool* symtab = upb_FileDef_Pool(file);
74 const upb_ExtensionRegistry* reg = upb_DefPool_ExtensionRegistry(symtab);
75 int64_t number = PyLong_AsLong(arg);
76 const upb_MiniTableExtension* ext =
77 (upb_MiniTableExtension*)upb_ExtensionRegistry_Lookup(reg, l, number);
78 if (ext) {
79 const upb_FieldDef* f = upb_DefPool_FindExtensionByMiniTable(symtab, ext);
80 return PyUpb_FieldDescriptor_Get(f);
81 } else {
82 Py_RETURN_NONE;
83 }
84 }
85
PyUpb_ExtensionDict_Dealloc(PyUpb_ExtensionDict * self)86 static void PyUpb_ExtensionDict_Dealloc(PyUpb_ExtensionDict* self) {
87 PyUpb_Message_ClearExtensionDict(self->msg);
88 Py_DECREF(self->msg);
89 PyUpb_Dealloc(self);
90 }
91
PyUpb_ExtensionDict_RichCompare(PyObject * _self,PyObject * _other,int opid)92 static PyObject* PyUpb_ExtensionDict_RichCompare(PyObject* _self,
93 PyObject* _other, int opid) {
94 // Only equality comparisons are implemented.
95 if (opid != Py_EQ && opid != Py_NE) {
96 Py_INCREF(Py_NotImplemented);
97 return Py_NotImplemented;
98 }
99 PyUpb_ExtensionDict* self = (PyUpb_ExtensionDict*)_self;
100 bool equals = false;
101 if (PyObject_TypeCheck(_other, Py_TYPE(_self))) {
102 PyUpb_ExtensionDict* other = (PyUpb_ExtensionDict*)_other;
103 equals = self->msg == other->msg;
104 }
105 bool ret = opid == Py_EQ ? equals : !equals;
106 return PyBool_FromLong(ret);
107 }
108
PyUpb_ExtensionDict_Contains(PyObject * _self,PyObject * key)109 static int PyUpb_ExtensionDict_Contains(PyObject* _self, PyObject* key) {
110 PyUpb_ExtensionDict* self = (PyUpb_ExtensionDict*)_self;
111 const upb_FieldDef* f = PyUpb_Message_GetExtensionDef(self->msg, key);
112 if (!f) return -1;
113 upb_Message* msg = PyUpb_Message_GetIfReified(self->msg);
114 if (!msg) return 0;
115 if (upb_FieldDef_IsRepeated(f)) {
116 upb_MessageValue val = upb_Message_GetFieldByDef(msg, f);
117 return upb_Array_Size(val.array_val) > 0;
118 } else {
119 return upb_Message_HasFieldByDef(msg, f);
120 }
121 }
122
PyUpb_ExtensionDict_Length(PyObject * _self)123 static Py_ssize_t PyUpb_ExtensionDict_Length(PyObject* _self) {
124 PyUpb_ExtensionDict* self = (PyUpb_ExtensionDict*)_self;
125 upb_Message* msg = PyUpb_Message_GetIfReified(self->msg);
126 return msg ? upb_Message_ExtensionCount(msg) : 0;
127 }
128
PyUpb_ExtensionDict_Subscript(PyObject * _self,PyObject * key)129 static PyObject* PyUpb_ExtensionDict_Subscript(PyObject* _self, PyObject* key) {
130 PyUpb_ExtensionDict* self = (PyUpb_ExtensionDict*)_self;
131 const upb_FieldDef* f = PyUpb_Message_GetExtensionDef(self->msg, key);
132 if (!f) return NULL;
133 return PyUpb_Message_GetFieldValue(self->msg, f);
134 }
135
PyUpb_ExtensionDict_AssignSubscript(PyObject * _self,PyObject * key,PyObject * val)136 static int PyUpb_ExtensionDict_AssignSubscript(PyObject* _self, PyObject* key,
137 PyObject* val) {
138 PyUpb_ExtensionDict* self = (PyUpb_ExtensionDict*)_self;
139 const upb_FieldDef* f = PyUpb_Message_GetExtensionDef(self->msg, key);
140 if (!f) return -1;
141 if (val) {
142 return PyUpb_Message_SetFieldValue(self->msg, f, val, PyExc_TypeError);
143 } else {
144 PyUpb_Message_DoClearField(self->msg, f);
145 return 0;
146 }
147 }
148
149 static PyObject* PyUpb_ExtensionIterator_New(PyObject* _ext_dict);
150
151 static PyMethodDef PyUpb_ExtensionDict_Methods[] = {
152 {"_FindExtensionByName", PyUpb_ExtensionDict_FindExtensionByName, METH_O,
153 "Finds an extension by name."},
154 {"_FindExtensionByNumber", PyUpb_ExtensionDict_FindExtensionByNumber,
155 METH_O, "Finds an extension by number."},
156 {NULL, NULL},
157 };
158
159 static PyType_Slot PyUpb_ExtensionDict_Slots[] = {
160 {Py_tp_dealloc, PyUpb_ExtensionDict_Dealloc},
161 {Py_tp_methods, PyUpb_ExtensionDict_Methods},
162 //{Py_tp_getset, PyUpb_ExtensionDict_Getters},
163 //{Py_tp_hash, PyObject_HashNotImplemented},
164 {Py_tp_richcompare, PyUpb_ExtensionDict_RichCompare},
165 {Py_tp_iter, PyUpb_ExtensionIterator_New},
166 {Py_sq_contains, PyUpb_ExtensionDict_Contains},
167 {Py_sq_length, PyUpb_ExtensionDict_Length},
168 {Py_mp_length, PyUpb_ExtensionDict_Length},
169 {Py_mp_subscript, PyUpb_ExtensionDict_Subscript},
170 {Py_mp_ass_subscript, PyUpb_ExtensionDict_AssignSubscript},
171 {0, NULL}};
172
173 static PyType_Spec PyUpb_ExtensionDict_Spec = {
174 PYUPB_MODULE_NAME ".ExtensionDict", // tp_name
175 sizeof(PyUpb_ExtensionDict), // tp_basicsize
176 0, // tp_itemsize
177 Py_TPFLAGS_DEFAULT, // tp_flags
178 PyUpb_ExtensionDict_Slots,
179 };
180
181 // -----------------------------------------------------------------------------
182 // ExtensionIterator
183 // -----------------------------------------------------------------------------
184
185 typedef struct {
186 PyObject_HEAD;
187 PyObject* msg;
188 size_t iter;
189 } PyUpb_ExtensionIterator;
190
PyUpb_ExtensionIterator_New(PyObject * _ext_dict)191 static PyObject* PyUpb_ExtensionIterator_New(PyObject* _ext_dict) {
192 PyUpb_ExtensionDict* ext_dict = (PyUpb_ExtensionDict*)_ext_dict;
193 PyUpb_ModuleState* state = PyUpb_ModuleState_Get();
194 PyUpb_ExtensionIterator* iter =
195 (void*)PyType_GenericAlloc(state->extension_iterator_type, 0);
196 if (!iter) return NULL;
197 iter->msg = ext_dict->msg;
198 iter->iter = kUpb_Message_Begin;
199 Py_INCREF(iter->msg);
200 return &iter->ob_base;
201 }
202
PyUpb_ExtensionIterator_Dealloc(void * _self)203 static void PyUpb_ExtensionIterator_Dealloc(void* _self) {
204 PyUpb_ExtensionIterator* self = (PyUpb_ExtensionIterator*)_self;
205 Py_DECREF(self->msg);
206 PyUpb_Dealloc(_self);
207 }
208
PyUpb_ExtensionIterator_IterNext(PyObject * _self)209 PyObject* PyUpb_ExtensionIterator_IterNext(PyObject* _self) {
210 PyUpb_ExtensionIterator* self = (PyUpb_ExtensionIterator*)_self;
211 upb_Message* msg = PyUpb_Message_GetIfReified(self->msg);
212 if (!msg) return NULL;
213 const upb_MessageDef* m = PyUpb_Message_GetMsgdef(self->msg);
214 const upb_DefPool* symtab = upb_FileDef_Pool(upb_MessageDef_File(m));
215 while (true) {
216 const upb_FieldDef* f;
217 upb_MessageValue val;
218 if (!upb_Message_Next(msg, m, symtab, &f, &val, &self->iter)) return NULL;
219 if (upb_FieldDef_IsExtension(f)) return PyUpb_FieldDescriptor_Get(f);
220 }
221 }
222
223 static PyType_Slot PyUpb_ExtensionIterator_Slots[] = {
224 {Py_tp_dealloc, PyUpb_ExtensionIterator_Dealloc},
225 {Py_tp_iter, PyObject_SelfIter},
226 {Py_tp_iternext, PyUpb_ExtensionIterator_IterNext},
227 {0, NULL}};
228
229 static PyType_Spec PyUpb_ExtensionIterator_Spec = {
230 PYUPB_MODULE_NAME ".ExtensionIterator", // tp_name
231 sizeof(PyUpb_ExtensionIterator), // tp_basicsize
232 0, // tp_itemsize
233 Py_TPFLAGS_DEFAULT, // tp_flags
234 PyUpb_ExtensionIterator_Slots,
235 };
236
237 // -----------------------------------------------------------------------------
238 // Top Level
239 // -----------------------------------------------------------------------------
240
PyUpb_InitExtensionDict(PyObject * m)241 bool PyUpb_InitExtensionDict(PyObject* m) {
242 PyUpb_ModuleState* s = PyUpb_ModuleState_GetFromModule(m);
243
244 s->extension_dict_type = PyUpb_AddClass(m, &PyUpb_ExtensionDict_Spec);
245 s->extension_iterator_type = PyUpb_AddClass(m, &PyUpb_ExtensionIterator_Spec);
246
247 return s->extension_dict_type && s->extension_iterator_type;
248 }
249