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/descriptor_pool.h"
9
10 #include "google/protobuf/descriptor.upbdefs.h"
11 #include "python/convert.h"
12 #include "python/descriptor.h"
13 #include "python/message.h"
14 #include "python/protobuf.h"
15 #include "upb/base/upcast.h"
16 #include "upb/message/compare.h"
17 #include "upb/reflection/def.h"
18 #include "upb/util/def_to_proto.h"
19
20 // -----------------------------------------------------------------------------
21 // DescriptorPool
22 // -----------------------------------------------------------------------------
23
24 typedef struct {
25 PyObject_HEAD;
26 upb_DefPool* symtab;
27 PyObject* db; // The DescriptorDatabase underlying this pool. May be NULL.
28 } PyUpb_DescriptorPool;
29
PyUpb_DescriptorPool_GetDefaultPool(void)30 PyObject* PyUpb_DescriptorPool_GetDefaultPool(void) {
31 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
32 return s->default_pool;
33 }
34
PyUpb_DescriptorPool_GetFileProtoDef(void)35 const upb_MessageDef* PyUpb_DescriptorPool_GetFileProtoDef(void) {
36 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
37 if (!s->c_descriptor_symtab) {
38 s->c_descriptor_symtab = upb_DefPool_New();
39 }
40 return google_protobuf_FileDescriptorProto_getmsgdef(s->c_descriptor_symtab);
41 }
42
PyUpb_DescriptorPool_DoCreateWithCache(PyTypeObject * type,PyObject * db,PyUpb_WeakMap * obj_cache)43 static PyObject* PyUpb_DescriptorPool_DoCreateWithCache(
44 PyTypeObject* type, PyObject* db, PyUpb_WeakMap* obj_cache) {
45 PyUpb_DescriptorPool* pool = (void*)PyType_GenericAlloc(type, 0);
46 pool->symtab = upb_DefPool_New();
47 pool->db = db;
48 Py_XINCREF(pool->db);
49 PyUpb_WeakMap_Add(obj_cache, pool->symtab, &pool->ob_base);
50 return &pool->ob_base;
51 }
52
PyUpb_DescriptorPool_DoCreate(PyTypeObject * type,PyObject * db)53 static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type,
54 PyObject* db) {
55 return PyUpb_DescriptorPool_DoCreateWithCache(type, db,
56 PyUpb_ObjCache_Instance());
57 }
58
PyUpb_DescriptorPool_GetSymtab(PyObject * pool)59 upb_DefPool* PyUpb_DescriptorPool_GetSymtab(PyObject* pool) {
60 return ((PyUpb_DescriptorPool*)pool)->symtab;
61 }
62
PyUpb_DescriptorPool_Traverse(PyUpb_DescriptorPool * self,visitproc visit,void * arg)63 static int PyUpb_DescriptorPool_Traverse(PyUpb_DescriptorPool* self,
64 visitproc visit, void* arg) {
65 Py_VISIT(self->db);
66 return 0;
67 }
68
PyUpb_DescriptorPool_Clear(PyUpb_DescriptorPool * self)69 static int PyUpb_DescriptorPool_Clear(PyUpb_DescriptorPool* self) {
70 Py_CLEAR(self->db);
71 return 0;
72 }
73
PyUpb_DescriptorPool_Get(const upb_DefPool * symtab)74 PyObject* PyUpb_DescriptorPool_Get(const upb_DefPool* symtab) {
75 PyObject* pool = PyUpb_ObjCache_Get(symtab);
76 assert(pool);
77 return pool;
78 }
79
PyUpb_DescriptorPool_Dealloc(PyUpb_DescriptorPool * self)80 static void PyUpb_DescriptorPool_Dealloc(PyUpb_DescriptorPool* self) {
81 PyObject_GC_UnTrack(self);
82 PyUpb_DescriptorPool_Clear(self);
83 upb_DefPool_Free(self->symtab);
84 PyUpb_ObjCache_Delete(self->symtab);
85 PyUpb_Dealloc(self);
86 }
87
88 /*
89 * DescriptorPool.__new__()
90 *
91 * Implements:
92 * DescriptorPool(descriptor_db=None)
93 */
PyUpb_DescriptorPool_New(PyTypeObject * type,PyObject * args,PyObject * kwargs)94 static PyObject* PyUpb_DescriptorPool_New(PyTypeObject* type, PyObject* args,
95 PyObject* kwargs) {
96 char* kwlist[] = {"descriptor_db", 0};
97 PyObject* db = NULL;
98
99 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O", kwlist, &db)) {
100 return NULL;
101 }
102
103 if (db == Py_None) db = NULL;
104 return PyUpb_DescriptorPool_DoCreate(type, db);
105 }
106
107 static PyObject* PyUpb_DescriptorPool_DoAdd(PyObject* _self,
108 PyObject* file_desc);
109
PyUpb_DescriptorPool_TryLoadFileProto(PyUpb_DescriptorPool * self,PyObject * proto)110 static bool PyUpb_DescriptorPool_TryLoadFileProto(PyUpb_DescriptorPool* self,
111 PyObject* proto) {
112 if (proto == NULL) {
113 if (PyErr_ExceptionMatches(PyExc_KeyError)) {
114 // Expected error: item was simply not found.
115 PyErr_Clear();
116 return true; // We didn't accomplish our goal, but we didn't error out.
117 }
118 return false;
119 }
120 if (proto == Py_None) return true;
121 PyObject* ret = PyUpb_DescriptorPool_DoAdd((PyObject*)self, proto);
122 bool ok = ret != NULL;
123 Py_XDECREF(ret);
124 return ok;
125 }
126
PyUpb_DescriptorPool_TryLoadSymbol(PyUpb_DescriptorPool * self,PyObject * sym)127 static bool PyUpb_DescriptorPool_TryLoadSymbol(PyUpb_DescriptorPool* self,
128 PyObject* sym) {
129 if (!self->db) return false;
130 PyObject* file_proto =
131 PyObject_CallMethod(self->db, "FindFileContainingSymbol", "O", sym);
132 bool ret = PyUpb_DescriptorPool_TryLoadFileProto(self, file_proto);
133 Py_XDECREF(file_proto);
134 return ret;
135 }
136
PyUpb_DescriptorPool_TryLoadFilename(PyUpb_DescriptorPool * self,PyObject * filename)137 static bool PyUpb_DescriptorPool_TryLoadFilename(PyUpb_DescriptorPool* self,
138 PyObject* filename) {
139 if (!self->db) return false;
140 PyObject* file_proto =
141 PyObject_CallMethod(self->db, "FindFileByName", "O", filename);
142 bool ret = PyUpb_DescriptorPool_TryLoadFileProto(self, file_proto);
143 Py_XDECREF(file_proto);
144 return ret;
145 }
146
PyUpb_DescriptorPool_CheckNoDatabase(PyObject * _self)147 bool PyUpb_DescriptorPool_CheckNoDatabase(PyObject* _self) { return true; }
148
PyUpb_DescriptorPool_LoadDependentFiles(PyUpb_DescriptorPool * self,google_protobuf_FileDescriptorProto * proto)149 static bool PyUpb_DescriptorPool_LoadDependentFiles(
150 PyUpb_DescriptorPool* self, google_protobuf_FileDescriptorProto* proto) {
151 size_t n;
152 const upb_StringView* deps = google_protobuf_FileDescriptorProto_dependency(proto, &n);
153 for (size_t i = 0; i < n; i++) {
154 const upb_FileDef* dep = upb_DefPool_FindFileByNameWithSize(
155 self->symtab, deps[i].data, deps[i].size);
156 if (!dep) {
157 PyObject* filename =
158 PyUnicode_FromStringAndSize(deps[i].data, deps[i].size);
159 if (!filename) return false;
160 bool ok = PyUpb_DescriptorPool_TryLoadFilename(self, filename);
161 Py_DECREF(filename);
162 if (!ok) return false;
163 }
164 }
165 return true;
166 }
167
PyUpb_DescriptorPool_DoAddSerializedFile(PyObject * _self,PyObject * serialized_pb)168 static PyObject* PyUpb_DescriptorPool_DoAddSerializedFile(
169 PyObject* _self, PyObject* serialized_pb) {
170 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
171 upb_Arena* arena = upb_Arena_New();
172 if (!arena) PYUPB_RETURN_OOM;
173 PyObject* result = NULL;
174
175 char* buf;
176 Py_ssize_t size;
177 if (PyBytes_AsStringAndSize(serialized_pb, &buf, &size) < 0) {
178 goto done;
179 }
180
181 google_protobuf_FileDescriptorProto* proto =
182 google_protobuf_FileDescriptorProto_parse(buf, size, arena);
183 if (!proto) {
184 PyErr_SetString(PyExc_TypeError, "Couldn't parse file content!");
185 goto done;
186 }
187
188 upb_StringView name = google_protobuf_FileDescriptorProto_name(proto);
189 const upb_FileDef* file =
190 upb_DefPool_FindFileByNameWithSize(self->symtab, name.data, name.size);
191
192 if (file) {
193 // If the existing file is equal to the new file, then silently ignore the
194 // duplicate add.
195 google_protobuf_FileDescriptorProto* existing = upb_FileDef_ToProto(file, arena);
196 if (!existing) {
197 PyErr_SetNone(PyExc_MemoryError);
198 goto done;
199 }
200 const upb_MessageDef* m = PyUpb_DescriptorPool_GetFileProtoDef();
201 const int options = kUpb_CompareOption_IncludeUnknownFields;
202 if (upb_Message_IsEqualByDef(UPB_UPCAST(proto), UPB_UPCAST(existing), m,
203 options)) {
204 result = PyUpb_FileDescriptor_Get(file);
205 goto done;
206 }
207 }
208
209 if (self->db) {
210 if (!PyUpb_DescriptorPool_LoadDependentFiles(self, proto)) goto done;
211 }
212
213 upb_Status status;
214 upb_Status_Clear(&status);
215
216 const upb_FileDef* filedef =
217 upb_DefPool_AddFile(self->symtab, proto, &status);
218 if (!filedef) {
219 PyErr_Format(PyExc_TypeError,
220 "Couldn't build proto file into descriptor pool: %s",
221 upb_Status_ErrorMessage(&status));
222 goto done;
223 }
224
225 result = PyUpb_FileDescriptor_Get(filedef);
226
227 done:
228 upb_Arena_Free(arena);
229 return result;
230 }
231
PyUpb_DescriptorPool_DoAdd(PyObject * _self,PyObject * file_desc)232 static PyObject* PyUpb_DescriptorPool_DoAdd(PyObject* _self,
233 PyObject* file_desc) {
234 if (!PyUpb_Message_Verify(file_desc)) return NULL;
235 const upb_MessageDef* m = PyUpb_Message_GetMsgdef(file_desc);
236 const char* file_proto_name =
237 PYUPB_DESCRIPTOR_PROTO_PACKAGE ".FileDescriptorProto";
238 if (strcmp(upb_MessageDef_FullName(m), file_proto_name) != 0) {
239 return PyErr_Format(PyExc_TypeError, "Can only add FileDescriptorProto");
240 }
241 PyObject* subargs = PyTuple_New(0);
242 if (!subargs) return NULL;
243 PyObject* serialized =
244 PyUpb_Message_SerializeToString(file_desc, subargs, NULL);
245 Py_DECREF(subargs);
246 if (!serialized) return NULL;
247 PyObject* ret = PyUpb_DescriptorPool_DoAddSerializedFile(_self, serialized);
248 Py_DECREF(serialized);
249 return ret;
250 }
251
252 /*
253 * PyUpb_DescriptorPool_AddSerializedFile()
254 *
255 * Implements:
256 * DescriptorPool.AddSerializedFile(self, serialized_file_descriptor)
257 *
258 * Adds the given serialized FileDescriptorProto to the pool.
259 */
PyUpb_DescriptorPool_AddSerializedFile(PyObject * _self,PyObject * serialized_pb)260 static PyObject* PyUpb_DescriptorPool_AddSerializedFile(
261 PyObject* _self, PyObject* serialized_pb) {
262 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
263 if (self->db) {
264 PyErr_SetString(
265 PyExc_ValueError,
266 "Cannot call AddSerializedFile on a DescriptorPool that uses a "
267 "DescriptorDatabase. Add your file to the underlying database.");
268 return false;
269 }
270 return PyUpb_DescriptorPool_DoAddSerializedFile(_self, serialized_pb);
271 }
272
PyUpb_DescriptorPool_Add(PyObject * _self,PyObject * file_desc)273 static PyObject* PyUpb_DescriptorPool_Add(PyObject* _self,
274 PyObject* file_desc) {
275 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
276 if (self->db) {
277 PyErr_SetString(
278 PyExc_ValueError,
279 "Cannot call Add on a DescriptorPool that uses a DescriptorDatabase. "
280 "Add your file to the underlying database.");
281 return false;
282 }
283 return PyUpb_DescriptorPool_DoAdd(_self, file_desc);
284 }
285
PyUpb_DescriptorPool_SetFeatureSetDefaults(PyObject * _self,PyObject * defaults)286 static PyObject* PyUpb_DescriptorPool_SetFeatureSetDefaults(
287 PyObject* _self, PyObject* defaults) {
288 if (!PyUpb_Message_Verify(defaults)) {
289 return PyErr_Format(PyExc_TypeError,
290 "SetFeatureSetDefaults called with invalid type");
291 }
292 const upb_MessageDef* m = PyUpb_Message_GetMsgdef(defaults);
293 const char* file_proto_name =
294 PYUPB_DESCRIPTOR_PROTO_PACKAGE ".FeatureSetDefaults";
295 if (strcmp(upb_MessageDef_FullName(m), file_proto_name) != 0) {
296 return PyErr_Format(
297 PyExc_TypeError,
298 "SetFeatureSetDefaults called with invalid type: got %s, expected %s",
299 upb_MessageDef_FullName(m), file_proto_name);
300 }
301 PyObject* subargs = PyTuple_New(0);
302 if (!subargs) return NULL;
303 PyObject* py_serialized =
304 PyUpb_Message_SerializeToString(defaults, subargs, NULL);
305 Py_DECREF(subargs);
306 if (!py_serialized) return NULL;
307 char* serialized;
308 Py_ssize_t size;
309 if (PyBytes_AsStringAndSize(py_serialized, &serialized, &size) < 0) {
310 goto err;
311 }
312
313 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
314 upb_Status status;
315 if (!upb_DefPool_SetFeatureSetDefaults(self->symtab, serialized, size,
316 &status)) {
317 PyErr_SetString(PyExc_ValueError, upb_Status_ErrorMessage(&status));
318 goto err;
319 }
320
321 Py_DECREF(py_serialized);
322 Py_RETURN_NONE;
323 err:
324 Py_DECREF(py_serialized);
325 return NULL;
326 }
327
328 /*
329 * PyUpb_DescriptorPool_FindFileByName()
330 *
331 * Implements:
332 * DescriptorPool.FindFileByName(self, name)
333 */
PyUpb_DescriptorPool_FindFileByName(PyObject * _self,PyObject * arg)334 static PyObject* PyUpb_DescriptorPool_FindFileByName(PyObject* _self,
335 PyObject* arg) {
336 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
337
338 const char* name = PyUpb_VerifyStrData(arg);
339 if (!name) return NULL;
340
341 const upb_FileDef* file = upb_DefPool_FindFileByName(self->symtab, name);
342 if (file == NULL && self->db) {
343 if (!PyUpb_DescriptorPool_TryLoadFilename(self, arg)) return NULL;
344 file = upb_DefPool_FindFileByName(self->symtab, name);
345 }
346 if (file == NULL) {
347 return PyErr_Format(PyExc_KeyError, "Couldn't find file %.200s", name);
348 }
349
350 return PyUpb_FileDescriptor_Get(file);
351 }
352
353 /*
354 * PyUpb_DescriptorPool_FindExtensionByName()
355 *
356 * Implements:
357 * DescriptorPool.FindExtensionByName(self, name)
358 */
PyUpb_DescriptorPool_FindExtensionByName(PyObject * _self,PyObject * arg)359 static PyObject* PyUpb_DescriptorPool_FindExtensionByName(PyObject* _self,
360 PyObject* arg) {
361 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
362
363 const char* name = PyUpb_VerifyStrData(arg);
364 if (!name) return NULL;
365
366 const upb_FieldDef* field =
367 upb_DefPool_FindExtensionByName(self->symtab, name);
368 if (field == NULL && self->db) {
369 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
370 field = upb_DefPool_FindExtensionByName(self->symtab, name);
371 }
372 if (field == NULL) {
373 return PyErr_Format(PyExc_KeyError, "Couldn't find extension %.200s", name);
374 }
375
376 return PyUpb_FieldDescriptor_Get(field);
377 }
378
379 /*
380 * PyUpb_DescriptorPool_FindMessageTypeByName()
381 *
382 * Implements:
383 * DescriptorPool.FindMessageTypeByName(self, name)
384 */
PyUpb_DescriptorPool_FindMessageTypeByName(PyObject * _self,PyObject * arg)385 static PyObject* PyUpb_DescriptorPool_FindMessageTypeByName(PyObject* _self,
386 PyObject* arg) {
387 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
388
389 const char* name = PyUpb_VerifyStrData(arg);
390 if (!name) return NULL;
391
392 const upb_MessageDef* m = upb_DefPool_FindMessageByName(self->symtab, name);
393 if (m == NULL && self->db) {
394 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
395 m = upb_DefPool_FindMessageByName(self->symtab, name);
396 }
397 if (m == NULL) {
398 return PyErr_Format(PyExc_KeyError, "Couldn't find message %.200s", name);
399 }
400
401 return PyUpb_Descriptor_Get(m);
402 }
403
404 // Splits a dotted symbol like foo.bar.baz on the last dot. Returns the portion
405 // after the last dot (baz) and updates `*parent_size` to the length of the
406 // parent (foo.bar). Returns NULL if no dots were present.
PyUpb_DescriptorPool_SplitSymbolName(const char * sym,size_t * parent_size)407 static const char* PyUpb_DescriptorPool_SplitSymbolName(const char* sym,
408 size_t* parent_size) {
409 const char* last_dot = strrchr(sym, '.');
410 if (!last_dot) return NULL;
411 *parent_size = last_dot - sym;
412 return last_dot + 1;
413 }
414
415 /*
416 * PyUpb_DescriptorPool_FindFieldByName()
417 *
418 * Implements:
419 * DescriptorPool.FindFieldByName(self, name)
420 */
PyUpb_DescriptorPool_FindFieldByName(PyObject * _self,PyObject * arg)421 static PyObject* PyUpb_DescriptorPool_FindFieldByName(PyObject* _self,
422 PyObject* arg) {
423 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
424
425 const char* name = PyUpb_VerifyStrData(arg);
426 if (!name) return NULL;
427
428 size_t parent_size;
429 const char* child = PyUpb_DescriptorPool_SplitSymbolName(name, &parent_size);
430 const upb_FieldDef* f = NULL;
431 if (child) {
432 const upb_MessageDef* parent =
433 upb_DefPool_FindMessageByNameWithSize(self->symtab, name, parent_size);
434 if (parent == NULL && self->db) {
435 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
436 parent = upb_DefPool_FindMessageByNameWithSize(self->symtab, name,
437 parent_size);
438 }
439 if (parent) {
440 f = upb_MessageDef_FindFieldByName(parent, child);
441 }
442 }
443
444 if (!f) {
445 return PyErr_Format(PyExc_KeyError, "Couldn't find message %.200s", name);
446 }
447
448 return PyUpb_FieldDescriptor_Get(f);
449 }
450
451 /*
452 * PyUpb_DescriptorPool_FindEnumTypeByName()
453 *
454 * Implements:
455 * DescriptorPool.FindEnumTypeByName(self, name)
456 */
PyUpb_DescriptorPool_FindEnumTypeByName(PyObject * _self,PyObject * arg)457 static PyObject* PyUpb_DescriptorPool_FindEnumTypeByName(PyObject* _self,
458 PyObject* arg) {
459 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
460
461 const char* name = PyUpb_VerifyStrData(arg);
462 if (!name) return NULL;
463
464 const upb_EnumDef* e = upb_DefPool_FindEnumByName(self->symtab, name);
465 if (e == NULL && self->db) {
466 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
467 e = upb_DefPool_FindEnumByName(self->symtab, name);
468 }
469 if (e == NULL) {
470 return PyErr_Format(PyExc_KeyError, "Couldn't find enum %.200s", name);
471 }
472
473 return PyUpb_EnumDescriptor_Get(e);
474 }
475
476 /*
477 * PyUpb_DescriptorPool_FindOneofByName()
478 *
479 * Implements:
480 * DescriptorPool.FindOneofByName(self, name)
481 */
PyUpb_DescriptorPool_FindOneofByName(PyObject * _self,PyObject * arg)482 static PyObject* PyUpb_DescriptorPool_FindOneofByName(PyObject* _self,
483 PyObject* arg) {
484 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
485
486 const char* name = PyUpb_VerifyStrData(arg);
487 if (!name) return NULL;
488
489 size_t parent_size;
490 const char* child = PyUpb_DescriptorPool_SplitSymbolName(name, &parent_size);
491
492 if (child) {
493 const upb_MessageDef* parent =
494 upb_DefPool_FindMessageByNameWithSize(self->symtab, name, parent_size);
495 if (parent == NULL && self->db) {
496 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
497 parent = upb_DefPool_FindMessageByNameWithSize(self->symtab, name,
498 parent_size);
499 }
500 if (parent) {
501 const upb_OneofDef* o = upb_MessageDef_FindOneofByName(parent, child);
502 return PyUpb_OneofDescriptor_Get(o);
503 }
504 }
505
506 return PyErr_Format(PyExc_KeyError, "Couldn't find oneof %.200s", name);
507 }
508
PyUpb_DescriptorPool_FindServiceByName(PyObject * _self,PyObject * arg)509 static PyObject* PyUpb_DescriptorPool_FindServiceByName(PyObject* _self,
510 PyObject* arg) {
511 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
512
513 const char* name = PyUpb_VerifyStrData(arg);
514 if (!name) return NULL;
515
516 const upb_ServiceDef* s = upb_DefPool_FindServiceByName(self->symtab, name);
517 if (s == NULL && self->db) {
518 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
519 s = upb_DefPool_FindServiceByName(self->symtab, name);
520 }
521 if (s == NULL) {
522 return PyErr_Format(PyExc_KeyError, "Couldn't find service %.200s", name);
523 }
524
525 return PyUpb_ServiceDescriptor_Get(s);
526 }
527
PyUpb_DescriptorPool_FindMethodByName(PyObject * _self,PyObject * arg)528 static PyObject* PyUpb_DescriptorPool_FindMethodByName(PyObject* _self,
529 PyObject* arg) {
530 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
531
532 const char* name = PyUpb_VerifyStrData(arg);
533 if (!name) return NULL;
534 size_t parent_size;
535 const char* child = PyUpb_DescriptorPool_SplitSymbolName(name, &parent_size);
536
537 if (!child) goto err;
538 const upb_ServiceDef* parent =
539 upb_DefPool_FindServiceByNameWithSize(self->symtab, name, parent_size);
540 if (parent == NULL && self->db) {
541 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
542 parent =
543 upb_DefPool_FindServiceByNameWithSize(self->symtab, name, parent_size);
544 }
545 if (!parent) goto err;
546 const upb_MethodDef* m = upb_ServiceDef_FindMethodByName(parent, child);
547 if (!m) goto err;
548 return PyUpb_MethodDescriptor_Get(m);
549
550 err:
551 return PyErr_Format(PyExc_KeyError, "Couldn't find method %.200s", name);
552 }
553
PyUpb_DescriptorPool_FindFileContainingSymbol(PyObject * _self,PyObject * arg)554 static PyObject* PyUpb_DescriptorPool_FindFileContainingSymbol(PyObject* _self,
555 PyObject* arg) {
556 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
557
558 const char* name = PyUpb_VerifyStrData(arg);
559 if (!name) return NULL;
560
561 const upb_FileDef* f =
562 upb_DefPool_FindFileContainingSymbol(self->symtab, name);
563 if (f == NULL && self->db) {
564 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
565 f = upb_DefPool_FindFileContainingSymbol(self->symtab, name);
566 }
567 if (f == NULL) {
568 return PyErr_Format(PyExc_KeyError, "Couldn't find symbol %.200s", name);
569 }
570
571 return PyUpb_FileDescriptor_Get(f);
572 }
573
PyUpb_DescriptorPool_FindExtensionByNumber(PyObject * _self,PyObject * args)574 static PyObject* PyUpb_DescriptorPool_FindExtensionByNumber(PyObject* _self,
575 PyObject* args) {
576 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
577 PyObject* message_descriptor;
578 int number;
579 if (!PyArg_ParseTuple(args, "Oi", &message_descriptor, &number)) {
580 return NULL;
581 }
582
583 const upb_FieldDef* f = upb_DefPool_FindExtensionByNumber(
584 self->symtab, PyUpb_Descriptor_GetDef(message_descriptor), number);
585 if (f == NULL) {
586 return PyErr_Format(PyExc_KeyError, "Couldn't find Extension %d", number);
587 }
588
589 return PyUpb_FieldDescriptor_Get(f);
590 }
591
PyUpb_DescriptorPool_FindAllExtensions(PyObject * _self,PyObject * msg_desc)592 static PyObject* PyUpb_DescriptorPool_FindAllExtensions(PyObject* _self,
593 PyObject* msg_desc) {
594 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
595 const upb_MessageDef* m = PyUpb_Descriptor_GetDef(msg_desc);
596 size_t n;
597 const upb_FieldDef** ext = upb_DefPool_GetAllExtensions(self->symtab, m, &n);
598 PyObject* ret = PyList_New(n);
599 if (!ret) goto done;
600 for (size_t i = 0; i < n; i++) {
601 PyObject* field = PyUpb_FieldDescriptor_Get(ext[i]);
602 if (!field) {
603 Py_DECREF(ret);
604 ret = NULL;
605 goto done;
606 }
607 PyList_SetItem(ret, i, field);
608 }
609 done:
610 free(ext);
611 return ret;
612 }
613
614 static PyMethodDef PyUpb_DescriptorPool_Methods[] = {
615 {"Add", PyUpb_DescriptorPool_Add, METH_O,
616 "Adds the FileDescriptorProto and its types to this pool."},
617 {"AddSerializedFile", PyUpb_DescriptorPool_AddSerializedFile, METH_O,
618 "Adds a serialized FileDescriptorProto to this pool."},
619 {"SetFeatureSetDefaults", PyUpb_DescriptorPool_SetFeatureSetDefaults,
620 METH_O, "Sets the default feature mappings used during the build."},
621 {"FindFileByName", PyUpb_DescriptorPool_FindFileByName, METH_O,
622 "Searches for a file descriptor by its .proto name."},
623 {"FindMessageTypeByName", PyUpb_DescriptorPool_FindMessageTypeByName,
624 METH_O, "Searches for a message descriptor by full name."},
625 {"FindFieldByName", PyUpb_DescriptorPool_FindFieldByName, METH_O,
626 "Searches for a field descriptor by full name."},
627 {"FindExtensionByName", PyUpb_DescriptorPool_FindExtensionByName, METH_O,
628 "Searches for extension descriptor by full name."},
629 {"FindEnumTypeByName", PyUpb_DescriptorPool_FindEnumTypeByName, METH_O,
630 "Searches for enum type descriptor by full name."},
631 {"FindOneofByName", PyUpb_DescriptorPool_FindOneofByName, METH_O,
632 "Searches for oneof descriptor by full name."},
633 {"FindServiceByName", PyUpb_DescriptorPool_FindServiceByName, METH_O,
634 "Searches for service descriptor by full name."},
635 {"FindMethodByName", PyUpb_DescriptorPool_FindMethodByName, METH_O,
636 "Searches for method descriptor by full name."},
637 {"FindFileContainingSymbol", PyUpb_DescriptorPool_FindFileContainingSymbol,
638 METH_O, "Gets the FileDescriptor containing the specified symbol."},
639 {"FindExtensionByNumber", PyUpb_DescriptorPool_FindExtensionByNumber,
640 METH_VARARGS, "Gets the extension descriptor for the given number."},
641 {"FindAllExtensions", PyUpb_DescriptorPool_FindAllExtensions, METH_O,
642 "Gets all known extensions of the given message descriptor."},
643 {NULL}};
644
645 static PyType_Slot PyUpb_DescriptorPool_Slots[] = {
646 {Py_tp_clear, PyUpb_DescriptorPool_Clear},
647 {Py_tp_dealloc, PyUpb_DescriptorPool_Dealloc},
648 {Py_tp_methods, PyUpb_DescriptorPool_Methods},
649 {Py_tp_new, PyUpb_DescriptorPool_New},
650 {Py_tp_traverse, PyUpb_DescriptorPool_Traverse},
651 {0, NULL}};
652
653 static PyType_Spec PyUpb_DescriptorPool_Spec = {
654 PYUPB_MODULE_NAME ".DescriptorPool",
655 sizeof(PyUpb_DescriptorPool),
656 0, // tp_itemsize
657 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
658 PyUpb_DescriptorPool_Slots,
659 };
660
661 // -----------------------------------------------------------------------------
662 // Top Level
663 // -----------------------------------------------------------------------------
664
PyUpb_InitDescriptorPool(PyObject * m)665 bool PyUpb_InitDescriptorPool(PyObject* m) {
666 PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
667 PyTypeObject* descriptor_pool_type =
668 PyUpb_AddClass(m, &PyUpb_DescriptorPool_Spec);
669
670 if (!descriptor_pool_type) return false;
671
672 state->default_pool = PyUpb_DescriptorPool_DoCreateWithCache(
673 descriptor_pool_type, NULL, state->obj_cache);
674 return state->default_pool &&
675 PyModule_AddObject(m, "default_pool", state->default_pool) == 0;
676 }
677