店播爬取Python脚本

descriptor_pool.cc 28KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818
  1. // Protocol Buffers - Google's data interchange format
  2. // Copyright 2008 Google Inc. All rights reserved.
  3. // https://developers.google.com/protocol-buffers/
  4. //
  5. // Redistribution and use in source and binary forms, with or without
  6. // modification, are permitted provided that the following conditions are
  7. // met:
  8. //
  9. // * Redistributions of source code must retain the above copyright
  10. // notice, this list of conditions and the following disclaimer.
  11. // * Redistributions in binary form must reproduce the above
  12. // copyright notice, this list of conditions and the following disclaimer
  13. // in the documentation and/or other materials provided with the
  14. // distribution.
  15. // * Neither the name of Google Inc. nor the names of its
  16. // contributors may be used to endorse or promote products derived from
  17. // this software without specific prior written permission.
  18. //
  19. // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  20. // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  21. // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  22. // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  23. // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  24. // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  25. // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  26. // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  27. // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  28. // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  29. // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  30. // Implements the DescriptorPool, which collects all descriptors.
  31. #include <unordered_map>
  32. #include <Python.h>
  33. #include <google/protobuf/descriptor.pb.h>
  34. #include <google/protobuf/pyext/descriptor.h>
  35. #include <google/protobuf/pyext/descriptor_database.h>
  36. #include <google/protobuf/pyext/descriptor_pool.h>
  37. #include <google/protobuf/pyext/message.h>
  38. #include <google/protobuf/pyext/message_factory.h>
  39. #include <google/protobuf/pyext/scoped_pyobject_ptr.h>
  40. #include <google/protobuf/stubs/hash.h>
  41. #if PY_MAJOR_VERSION >= 3
  42. #define PyString_FromStringAndSize PyUnicode_FromStringAndSize
  43. #if PY_VERSION_HEX < 0x03030000
  44. #error "Python 3.0 - 3.2 are not supported."
  45. #endif
  46. #define PyString_AsStringAndSize(ob, charpp, sizep) \
  47. (PyUnicode_Check(ob) ? ((*(charpp) = const_cast<char*>( \
  48. PyUnicode_AsUTF8AndSize(ob, (sizep)))) == NULL \
  49. ? -1 \
  50. : 0) \
  51. : PyBytes_AsStringAndSize(ob, (charpp), (sizep)))
  52. #endif
  53. namespace google {
  54. namespace protobuf {
  55. namespace python {
  56. // A map to cache Python Pools per C++ pointer.
  57. // Pointers are not owned here, and belong to the PyDescriptorPool.
  58. static std::unordered_map<const DescriptorPool*, PyDescriptorPool*>*
  59. descriptor_pool_map;
  60. namespace cdescriptor_pool {
  61. // Collects errors that occur during proto file building to allow them to be
  62. // propagated in the python exception instead of only living in ERROR logs.
  63. class BuildFileErrorCollector : public DescriptorPool::ErrorCollector {
  64. public:
  65. BuildFileErrorCollector() : error_message(""), had_errors_(false) {}
  66. void AddError(const std::string& filename, const std::string& element_name,
  67. const Message* descriptor, ErrorLocation location,
  68. const std::string& message) override {
  69. // Replicates the logging behavior that happens in the C++ implementation
  70. // when an error collector is not passed in.
  71. if (!had_errors_) {
  72. error_message +=
  73. ("Invalid proto descriptor for file \"" + filename + "\":\n");
  74. had_errors_ = true;
  75. }
  76. // As this only happens on failure and will result in the program not
  77. // running at all, no effort is made to optimize this string manipulation.
  78. error_message += (" " + element_name + ": " + message + "\n");
  79. }
  80. void Clear() {
  81. had_errors_ = false;
  82. error_message = "";
  83. }
  84. std::string error_message;
  85. private:
  86. bool had_errors_;
  87. };
  88. // Create a Python DescriptorPool object, but does not fill the "pool"
  89. // attribute.
  90. static PyDescriptorPool* _CreateDescriptorPool() {
  91. PyDescriptorPool* cpool = PyObject_GC_New(
  92. PyDescriptorPool, &PyDescriptorPool_Type);
  93. if (cpool == NULL) {
  94. return NULL;
  95. }
  96. cpool->error_collector = nullptr;
  97. cpool->underlay = NULL;
  98. cpool->database = NULL;
  99. cpool->is_owned = false;
  100. cpool->is_mutable = false;
  101. cpool->descriptor_options = new std::unordered_map<const void*, PyObject*>();
  102. cpool->py_message_factory = message_factory::NewMessageFactory(
  103. &PyMessageFactory_Type, cpool);
  104. if (cpool->py_message_factory == NULL) {
  105. Py_DECREF(cpool);
  106. return NULL;
  107. }
  108. PyObject_GC_Track(cpool);
  109. return cpool;
  110. }
  111. // Create a Python DescriptorPool, using the given pool as an underlay:
  112. // new messages will be added to a custom pool, not to the underlay.
  113. //
  114. // Ownership of the underlay is not transferred, its pointer should
  115. // stay alive.
  116. static PyDescriptorPool* PyDescriptorPool_NewWithUnderlay(
  117. const DescriptorPool* underlay) {
  118. PyDescriptorPool* cpool = _CreateDescriptorPool();
  119. if (cpool == NULL) {
  120. return NULL;
  121. }
  122. cpool->pool = new DescriptorPool(underlay);
  123. cpool->is_owned = true;
  124. cpool->is_mutable = true;
  125. cpool->underlay = underlay;
  126. if (!descriptor_pool_map->insert(
  127. std::make_pair(cpool->pool, cpool)).second) {
  128. // Should never happen -- would indicate an internal error / bug.
  129. PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
  130. return NULL;
  131. }
  132. return cpool;
  133. }
  134. static PyDescriptorPool* PyDescriptorPool_NewWithDatabase(
  135. DescriptorDatabase* database) {
  136. PyDescriptorPool* cpool = _CreateDescriptorPool();
  137. if (cpool == NULL) {
  138. return NULL;
  139. }
  140. if (database != NULL) {
  141. cpool->error_collector = new BuildFileErrorCollector();
  142. cpool->pool = new DescriptorPool(database, cpool->error_collector);
  143. cpool->is_mutable = false;
  144. cpool->database = database;
  145. } else {
  146. cpool->pool = new DescriptorPool();
  147. cpool->is_mutable = true;
  148. }
  149. cpool->is_owned = true;
  150. if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool)).second) {
  151. // Should never happen -- would indicate an internal error / bug.
  152. PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
  153. return NULL;
  154. }
  155. return cpool;
  156. }
  157. // The public DescriptorPool constructor.
  158. static PyObject* New(PyTypeObject* type,
  159. PyObject* args, PyObject* kwargs) {
  160. static const char* kwlist[] = {"descriptor_db", 0};
  161. PyObject* py_database = NULL;
  162. if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O",
  163. const_cast<char**>(kwlist), &py_database)) {
  164. return NULL;
  165. }
  166. DescriptorDatabase* database = NULL;
  167. if (py_database && py_database != Py_None) {
  168. database = new PyDescriptorDatabase(py_database);
  169. }
  170. return reinterpret_cast<PyObject*>(
  171. PyDescriptorPool_NewWithDatabase(database));
  172. }
  173. static void Dealloc(PyObject* pself) {
  174. PyDescriptorPool* self = reinterpret_cast<PyDescriptorPool*>(pself);
  175. descriptor_pool_map->erase(self->pool);
  176. Py_CLEAR(self->py_message_factory);
  177. for (std::unordered_map<const void*, PyObject*>::iterator it =
  178. self->descriptor_options->begin();
  179. it != self->descriptor_options->end(); ++it) {
  180. Py_DECREF(it->second);
  181. }
  182. delete self->descriptor_options;
  183. delete self->database;
  184. if (self->is_owned) {
  185. delete self->pool;
  186. }
  187. delete self->error_collector;
  188. Py_TYPE(self)->tp_free(pself);
  189. }
  190. static int GcTraverse(PyObject* pself, visitproc visit, void* arg) {
  191. PyDescriptorPool* self = reinterpret_cast<PyDescriptorPool*>(pself);
  192. Py_VISIT(self->py_message_factory);
  193. return 0;
  194. }
  195. static int GcClear(PyObject* pself) {
  196. PyDescriptorPool* self = reinterpret_cast<PyDescriptorPool*>(pself);
  197. Py_CLEAR(self->py_message_factory);
  198. return 0;
  199. }
  200. PyObject* SetErrorFromCollector(DescriptorPool::ErrorCollector* self,
  201. const char* name, const char* error_type) {
  202. BuildFileErrorCollector* error_collector =
  203. reinterpret_cast<BuildFileErrorCollector*>(self);
  204. if (error_collector && !error_collector->error_message.empty()) {
  205. PyErr_Format(PyExc_KeyError, "Couldn't build file for %s %.200s\n%s",
  206. error_type, name, error_collector->error_message.c_str());
  207. error_collector->Clear();
  208. return NULL;
  209. }
  210. PyErr_Format(PyExc_KeyError, "Couldn't find %s %.200s", error_type, name);
  211. return NULL;
  212. }
  213. static PyObject* FindMessageByName(PyObject* self, PyObject* arg) {
  214. Py_ssize_t name_size;
  215. char* name;
  216. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  217. return NULL;
  218. }
  219. const Descriptor* message_descriptor =
  220. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindMessageTypeByName(
  221. StringParam(name, name_size));
  222. if (message_descriptor == NULL) {
  223. return SetErrorFromCollector(
  224. reinterpret_cast<PyDescriptorPool*>(self)->error_collector, name,
  225. "message");
  226. }
  227. return PyMessageDescriptor_FromDescriptor(message_descriptor);
  228. }
  229. static PyObject* FindFileByName(PyObject* self, PyObject* arg) {
  230. Py_ssize_t name_size;
  231. char* name;
  232. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  233. return NULL;
  234. }
  235. PyDescriptorPool* py_pool = reinterpret_cast<PyDescriptorPool*>(self);
  236. const FileDescriptor* file_descriptor =
  237. py_pool->pool->FindFileByName(StringParam(name, name_size));
  238. if (file_descriptor == NULL) {
  239. return SetErrorFromCollector(py_pool->error_collector, name, "file");
  240. }
  241. return PyFileDescriptor_FromDescriptor(file_descriptor);
  242. }
  243. PyObject* FindFieldByName(PyDescriptorPool* self, PyObject* arg) {
  244. Py_ssize_t name_size;
  245. char* name;
  246. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  247. return NULL;
  248. }
  249. const FieldDescriptor* field_descriptor =
  250. self->pool->FindFieldByName(StringParam(name, name_size));
  251. if (field_descriptor == NULL) {
  252. return SetErrorFromCollector(self->error_collector, name, "field");
  253. }
  254. return PyFieldDescriptor_FromDescriptor(field_descriptor);
  255. }
  256. static PyObject* FindFieldByNameMethod(PyObject* self, PyObject* arg) {
  257. return FindFieldByName(reinterpret_cast<PyDescriptorPool*>(self), arg);
  258. }
  259. PyObject* FindExtensionByName(PyDescriptorPool* self, PyObject* arg) {
  260. Py_ssize_t name_size;
  261. char* name;
  262. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  263. return NULL;
  264. }
  265. const FieldDescriptor* field_descriptor =
  266. self->pool->FindExtensionByName(StringParam(name, name_size));
  267. if (field_descriptor == NULL) {
  268. return SetErrorFromCollector(self->error_collector, name,
  269. "extension field");
  270. }
  271. return PyFieldDescriptor_FromDescriptor(field_descriptor);
  272. }
  273. static PyObject* FindExtensionByNameMethod(PyObject* self, PyObject* arg) {
  274. return FindExtensionByName(reinterpret_cast<PyDescriptorPool*>(self), arg);
  275. }
  276. PyObject* FindEnumTypeByName(PyDescriptorPool* self, PyObject* arg) {
  277. Py_ssize_t name_size;
  278. char* name;
  279. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  280. return NULL;
  281. }
  282. const EnumDescriptor* enum_descriptor =
  283. self->pool->FindEnumTypeByName(StringParam(name, name_size));
  284. if (enum_descriptor == NULL) {
  285. return SetErrorFromCollector(self->error_collector, name, "enum");
  286. }
  287. return PyEnumDescriptor_FromDescriptor(enum_descriptor);
  288. }
  289. static PyObject* FindEnumTypeByNameMethod(PyObject* self, PyObject* arg) {
  290. return FindEnumTypeByName(reinterpret_cast<PyDescriptorPool*>(self), arg);
  291. }
  292. PyObject* FindOneofByName(PyDescriptorPool* self, PyObject* arg) {
  293. Py_ssize_t name_size;
  294. char* name;
  295. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  296. return NULL;
  297. }
  298. const OneofDescriptor* oneof_descriptor =
  299. self->pool->FindOneofByName(StringParam(name, name_size));
  300. if (oneof_descriptor == NULL) {
  301. return SetErrorFromCollector(self->error_collector, name, "oneof");
  302. }
  303. return PyOneofDescriptor_FromDescriptor(oneof_descriptor);
  304. }
  305. static PyObject* FindOneofByNameMethod(PyObject* self, PyObject* arg) {
  306. return FindOneofByName(reinterpret_cast<PyDescriptorPool*>(self), arg);
  307. }
  308. static PyObject* FindServiceByName(PyObject* self, PyObject* arg) {
  309. Py_ssize_t name_size;
  310. char* name;
  311. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  312. return NULL;
  313. }
  314. const ServiceDescriptor* service_descriptor =
  315. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindServiceByName(
  316. StringParam(name, name_size));
  317. if (service_descriptor == NULL) {
  318. return SetErrorFromCollector(
  319. reinterpret_cast<PyDescriptorPool*>(self)->error_collector, name,
  320. "service");
  321. }
  322. return PyServiceDescriptor_FromDescriptor(service_descriptor);
  323. }
  324. static PyObject* FindMethodByName(PyObject* self, PyObject* arg) {
  325. Py_ssize_t name_size;
  326. char* name;
  327. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  328. return NULL;
  329. }
  330. const MethodDescriptor* method_descriptor =
  331. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindMethodByName(
  332. StringParam(name, name_size));
  333. if (method_descriptor == NULL) {
  334. return SetErrorFromCollector(
  335. reinterpret_cast<PyDescriptorPool*>(self)->error_collector, name,
  336. "method");
  337. }
  338. return PyMethodDescriptor_FromDescriptor(method_descriptor);
  339. }
  340. static PyObject* FindFileContainingSymbol(PyObject* self, PyObject* arg) {
  341. Py_ssize_t name_size;
  342. char* name;
  343. if (PyString_AsStringAndSize(arg, &name, &name_size) < 0) {
  344. return NULL;
  345. }
  346. const FileDescriptor* file_descriptor =
  347. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindFileContainingSymbol(
  348. StringParam(name, name_size));
  349. if (file_descriptor == NULL) {
  350. return SetErrorFromCollector(
  351. reinterpret_cast<PyDescriptorPool*>(self)->error_collector, name,
  352. "symbol");
  353. }
  354. return PyFileDescriptor_FromDescriptor(file_descriptor);
  355. }
  356. static PyObject* FindExtensionByNumber(PyObject* self, PyObject* args) {
  357. PyObject* message_descriptor;
  358. int number;
  359. if (!PyArg_ParseTuple(args, "Oi", &message_descriptor, &number)) {
  360. return NULL;
  361. }
  362. const Descriptor* descriptor = PyMessageDescriptor_AsDescriptor(
  363. message_descriptor);
  364. if (descriptor == NULL) {
  365. return NULL;
  366. }
  367. const FieldDescriptor* extension_descriptor =
  368. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindExtensionByNumber(
  369. descriptor, number);
  370. if (extension_descriptor == NULL) {
  371. BuildFileErrorCollector* error_collector =
  372. reinterpret_cast<BuildFileErrorCollector*>(
  373. reinterpret_cast<PyDescriptorPool*>(self)->error_collector);
  374. if (error_collector && !error_collector->error_message.empty()) {
  375. PyErr_Format(PyExc_KeyError, "Couldn't build file for Extension %.d\n%s",
  376. number, error_collector->error_message.c_str());
  377. error_collector->Clear();
  378. return NULL;
  379. }
  380. PyErr_Format(PyExc_KeyError, "Couldn't find Extension %d", number);
  381. return NULL;
  382. }
  383. return PyFieldDescriptor_FromDescriptor(extension_descriptor);
  384. }
  385. static PyObject* FindAllExtensions(PyObject* self, PyObject* arg) {
  386. const Descriptor* descriptor = PyMessageDescriptor_AsDescriptor(arg);
  387. if (descriptor == NULL) {
  388. return NULL;
  389. }
  390. std::vector<const FieldDescriptor*> extensions;
  391. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindAllExtensions(
  392. descriptor, &extensions);
  393. ScopedPyObjectPtr result(PyList_New(extensions.size()));
  394. if (result == NULL) {
  395. return NULL;
  396. }
  397. for (int i = 0; i < extensions.size(); i++) {
  398. PyObject* extension = PyFieldDescriptor_FromDescriptor(extensions[i]);
  399. if (extension == NULL) {
  400. return NULL;
  401. }
  402. PyList_SET_ITEM(result.get(), i, extension); // Steals the reference.
  403. }
  404. return result.release();
  405. }
  406. // These functions should not exist -- the only valid way to create
  407. // descriptors is to call Add() or AddSerializedFile().
  408. // But these AddDescriptor() functions were created in Python and some people
  409. // call them, so we support them for now for compatibility.
  410. // However we do check that the existing descriptor already exists in the pool,
  411. // which appears to always be true for existing calls -- but then why do people
  412. // call a function that will just be a no-op?
  413. // TODO(amauryfa): Need to investigate further.
  414. static PyObject* AddFileDescriptor(PyObject* self, PyObject* descriptor) {
  415. const FileDescriptor* file_descriptor =
  416. PyFileDescriptor_AsDescriptor(descriptor);
  417. if (!file_descriptor) {
  418. return NULL;
  419. }
  420. if (file_descriptor !=
  421. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindFileByName(
  422. file_descriptor->name())) {
  423. PyErr_Format(PyExc_ValueError,
  424. "The file descriptor %s does not belong to this pool",
  425. file_descriptor->name().c_str());
  426. return NULL;
  427. }
  428. Py_RETURN_NONE;
  429. }
  430. static PyObject* AddDescriptor(PyObject* self, PyObject* descriptor) {
  431. const Descriptor* message_descriptor =
  432. PyMessageDescriptor_AsDescriptor(descriptor);
  433. if (!message_descriptor) {
  434. return NULL;
  435. }
  436. if (message_descriptor !=
  437. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindMessageTypeByName(
  438. message_descriptor->full_name())) {
  439. PyErr_Format(PyExc_ValueError,
  440. "The message descriptor %s does not belong to this pool",
  441. message_descriptor->full_name().c_str());
  442. return NULL;
  443. }
  444. Py_RETURN_NONE;
  445. }
  446. static PyObject* AddEnumDescriptor(PyObject* self, PyObject* descriptor) {
  447. const EnumDescriptor* enum_descriptor =
  448. PyEnumDescriptor_AsDescriptor(descriptor);
  449. if (!enum_descriptor) {
  450. return NULL;
  451. }
  452. if (enum_descriptor !=
  453. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindEnumTypeByName(
  454. enum_descriptor->full_name())) {
  455. PyErr_Format(PyExc_ValueError,
  456. "The enum descriptor %s does not belong to this pool",
  457. enum_descriptor->full_name().c_str());
  458. return NULL;
  459. }
  460. Py_RETURN_NONE;
  461. }
  462. static PyObject* AddExtensionDescriptor(PyObject* self, PyObject* descriptor) {
  463. const FieldDescriptor* extension_descriptor =
  464. PyFieldDescriptor_AsDescriptor(descriptor);
  465. if (!extension_descriptor) {
  466. return NULL;
  467. }
  468. if (extension_descriptor !=
  469. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindExtensionByName(
  470. extension_descriptor->full_name())) {
  471. PyErr_Format(PyExc_ValueError,
  472. "The extension descriptor %s does not belong to this pool",
  473. extension_descriptor->full_name().c_str());
  474. return NULL;
  475. }
  476. Py_RETURN_NONE;
  477. }
  478. static PyObject* AddServiceDescriptor(PyObject* self, PyObject* descriptor) {
  479. const ServiceDescriptor* service_descriptor =
  480. PyServiceDescriptor_AsDescriptor(descriptor);
  481. if (!service_descriptor) {
  482. return NULL;
  483. }
  484. if (service_descriptor !=
  485. reinterpret_cast<PyDescriptorPool*>(self)->pool->FindServiceByName(
  486. service_descriptor->full_name())) {
  487. PyErr_Format(PyExc_ValueError,
  488. "The service descriptor %s does not belong to this pool",
  489. service_descriptor->full_name().c_str());
  490. return NULL;
  491. }
  492. Py_RETURN_NONE;
  493. }
  494. // The code below loads new Descriptors from a serialized FileDescriptorProto.
  495. static PyObject* AddSerializedFile(PyObject* pself, PyObject* serialized_pb) {
  496. PyDescriptorPool* self = reinterpret_cast<PyDescriptorPool*>(pself);
  497. char* message_type;
  498. Py_ssize_t message_len;
  499. if (self->database != NULL) {
  500. PyErr_SetString(
  501. PyExc_ValueError,
  502. "Cannot call Add on a DescriptorPool that uses a DescriptorDatabase. "
  503. "Add your file to the underlying database.");
  504. return NULL;
  505. }
  506. if (!self->is_mutable) {
  507. PyErr_SetString(
  508. PyExc_ValueError,
  509. "This DescriptorPool is not mutable and cannot add new definitions.");
  510. return nullptr;
  511. }
  512. if (PyBytes_AsStringAndSize(serialized_pb, &message_type, &message_len) < 0) {
  513. return NULL;
  514. }
  515. FileDescriptorProto file_proto;
  516. if (!file_proto.ParseFromArray(message_type, message_len)) {
  517. PyErr_SetString(PyExc_TypeError, "Couldn't parse file content!");
  518. return NULL;
  519. }
  520. // If the file was already part of a C++ library, all its descriptors are in
  521. // the underlying pool. No need to do anything else.
  522. const FileDescriptor* generated_file = NULL;
  523. if (self->underlay) {
  524. generated_file = self->underlay->FindFileByName(file_proto.name());
  525. }
  526. if (generated_file != NULL) {
  527. return PyFileDescriptor_FromDescriptorWithSerializedPb(
  528. generated_file, serialized_pb);
  529. }
  530. BuildFileErrorCollector error_collector;
  531. const FileDescriptor* descriptor =
  532. // Pool is mutable, we can remove the "const".
  533. const_cast<DescriptorPool*>(self->pool)
  534. ->BuildFileCollectingErrors(file_proto, &error_collector);
  535. if (descriptor == NULL) {
  536. PyErr_Format(PyExc_TypeError,
  537. "Couldn't build proto file into descriptor pool!\n%s",
  538. error_collector.error_message.c_str());
  539. return NULL;
  540. }
  541. return PyFileDescriptor_FromDescriptorWithSerializedPb(
  542. descriptor, serialized_pb);
  543. }
  544. static PyObject* Add(PyObject* self, PyObject* file_descriptor_proto) {
  545. ScopedPyObjectPtr serialized_pb(
  546. PyObject_CallMethod(file_descriptor_proto, "SerializeToString", NULL));
  547. if (serialized_pb == NULL) {
  548. return NULL;
  549. }
  550. return AddSerializedFile(self, serialized_pb.get());
  551. }
  552. static PyMethodDef Methods[] = {
  553. { "Add", Add, METH_O,
  554. "Adds the FileDescriptorProto and its types to this pool." },
  555. { "AddSerializedFile", AddSerializedFile, METH_O,
  556. "Adds a serialized FileDescriptorProto to this pool." },
  557. // TODO(amauryfa): Understand why the Python implementation differs from
  558. // this one, ask users to use another API and deprecate these functions.
  559. { "AddFileDescriptor", AddFileDescriptor, METH_O,
  560. "No-op. Add() must have been called before." },
  561. { "AddDescriptor", AddDescriptor, METH_O,
  562. "No-op. Add() must have been called before." },
  563. { "AddEnumDescriptor", AddEnumDescriptor, METH_O,
  564. "No-op. Add() must have been called before." },
  565. { "AddExtensionDescriptor", AddExtensionDescriptor, METH_O,
  566. "No-op. Add() must have been called before." },
  567. { "AddServiceDescriptor", AddServiceDescriptor, METH_O,
  568. "No-op. Add() must have been called before." },
  569. { "FindFileByName", FindFileByName, METH_O,
  570. "Searches for a file descriptor by its .proto name." },
  571. { "FindMessageTypeByName", FindMessageByName, METH_O,
  572. "Searches for a message descriptor by full name." },
  573. { "FindFieldByName", FindFieldByNameMethod, METH_O,
  574. "Searches for a field descriptor by full name." },
  575. { "FindExtensionByName", FindExtensionByNameMethod, METH_O,
  576. "Searches for extension descriptor by full name." },
  577. { "FindEnumTypeByName", FindEnumTypeByNameMethod, METH_O,
  578. "Searches for enum type descriptor by full name." },
  579. { "FindOneofByName", FindOneofByNameMethod, METH_O,
  580. "Searches for oneof descriptor by full name." },
  581. { "FindServiceByName", FindServiceByName, METH_O,
  582. "Searches for service descriptor by full name." },
  583. { "FindMethodByName", FindMethodByName, METH_O,
  584. "Searches for method descriptor by full name." },
  585. { "FindFileContainingSymbol", FindFileContainingSymbol, METH_O,
  586. "Gets the FileDescriptor containing the specified symbol." },
  587. { "FindExtensionByNumber", FindExtensionByNumber, METH_VARARGS,
  588. "Gets the extension descriptor for the given number." },
  589. { "FindAllExtensions", FindAllExtensions, METH_O,
  590. "Gets all known extensions of the given message descriptor." },
  591. {NULL}
  592. };
  593. } // namespace cdescriptor_pool
  594. PyTypeObject PyDescriptorPool_Type = {
  595. PyVarObject_HEAD_INIT(&PyType_Type, 0) FULL_MODULE_NAME
  596. ".DescriptorPool", // tp_name
  597. sizeof(PyDescriptorPool), // tp_basicsize
  598. 0, // tp_itemsize
  599. cdescriptor_pool::Dealloc, // tp_dealloc
  600. 0, // tp_print
  601. 0, // tp_getattr
  602. 0, // tp_setattr
  603. 0, // tp_compare
  604. 0, // tp_repr
  605. 0, // tp_as_number
  606. 0, // tp_as_sequence
  607. 0, // tp_as_mapping
  608. 0, // tp_hash
  609. 0, // tp_call
  610. 0, // tp_str
  611. 0, // tp_getattro
  612. 0, // tp_setattro
  613. 0, // tp_as_buffer
  614. Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, // tp_flags
  615. "A Descriptor Pool", // tp_doc
  616. cdescriptor_pool::GcTraverse, // tp_traverse
  617. cdescriptor_pool::GcClear, // tp_clear
  618. 0, // tp_richcompare
  619. 0, // tp_weaklistoffset
  620. 0, // tp_iter
  621. 0, // tp_iternext
  622. cdescriptor_pool::Methods, // tp_methods
  623. 0, // tp_members
  624. 0, // tp_getset
  625. 0, // tp_base
  626. 0, // tp_dict
  627. 0, // tp_descr_get
  628. 0, // tp_descr_set
  629. 0, // tp_dictoffset
  630. 0, // tp_init
  631. 0, // tp_alloc
  632. cdescriptor_pool::New, // tp_new
  633. PyObject_GC_Del, // tp_free
  634. };
  635. // This is the DescriptorPool which contains all the definitions from the
  636. // generated _pb2.py modules.
  637. static PyDescriptorPool* python_generated_pool = NULL;
  638. bool InitDescriptorPool() {
  639. if (PyType_Ready(&PyDescriptorPool_Type) < 0)
  640. return false;
  641. // The Pool of messages declared in Python libraries.
  642. // generated_pool() contains all messages already linked in C++ libraries, and
  643. // is used as underlay.
  644. descriptor_pool_map =
  645. new std::unordered_map<const DescriptorPool*, PyDescriptorPool*>;
  646. python_generated_pool = cdescriptor_pool::PyDescriptorPool_NewWithUnderlay(
  647. DescriptorPool::generated_pool());
  648. if (python_generated_pool == NULL) {
  649. delete descriptor_pool_map;
  650. return false;
  651. }
  652. // Register this pool to be found for C++-generated descriptors.
  653. descriptor_pool_map->insert(
  654. std::make_pair(DescriptorPool::generated_pool(),
  655. python_generated_pool));
  656. return true;
  657. }
  658. // The default DescriptorPool used everywhere in this module.
  659. // Today it's the python_generated_pool.
  660. // TODO(amauryfa): Remove all usages of this function: the pool should be
  661. // derived from the context.
  662. PyDescriptorPool* GetDefaultDescriptorPool() {
  663. return python_generated_pool;
  664. }
  665. PyDescriptorPool* GetDescriptorPool_FromPool(const DescriptorPool* pool) {
  666. // Fast path for standard descriptors.
  667. if (pool == python_generated_pool->pool ||
  668. pool == DescriptorPool::generated_pool()) {
  669. return python_generated_pool;
  670. }
  671. std::unordered_map<const DescriptorPool*, PyDescriptorPool*>::iterator it =
  672. descriptor_pool_map->find(pool);
  673. if (it == descriptor_pool_map->end()) {
  674. PyErr_SetString(PyExc_KeyError, "Unknown descriptor pool");
  675. return NULL;
  676. }
  677. return it->second;
  678. }
  679. PyObject* PyDescriptorPool_FromPool(const DescriptorPool* pool) {
  680. PyDescriptorPool* existing_pool = GetDescriptorPool_FromPool(pool);
  681. if (existing_pool != nullptr) {
  682. Py_INCREF(existing_pool);
  683. return reinterpret_cast<PyObject*>(existing_pool);
  684. } else {
  685. PyErr_Clear();
  686. }
  687. PyDescriptorPool* cpool = cdescriptor_pool::_CreateDescriptorPool();
  688. if (cpool == nullptr) {
  689. return nullptr;
  690. }
  691. cpool->pool = const_cast<DescriptorPool*>(pool);
  692. cpool->is_owned = false;
  693. cpool->is_mutable = false;
  694. cpool->underlay = nullptr;
  695. if (!descriptor_pool_map->insert(std::make_pair(cpool->pool, cpool)).second) {
  696. // Should never happen -- We already checked the existence above.
  697. PyErr_SetString(PyExc_ValueError, "DescriptorPool already registered");
  698. return nullptr;
  699. }
  700. return reinterpret_cast<PyObject*>(cpool);
  701. }
  702. } // namespace python
  703. } // namespace protobuf
  704. } // namespace google