• 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/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