descriptor_pool.c 22 KB

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