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