blob: f6083f7d6bb4fa60da818d2d3d45b90c27ac1696 [file] [log] [blame]
Joshua Haberman94ece042021-08-18 12:38:26 -07001/*
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
28#include "python/descriptor_pool.h"
29
Joshua Haberman1c955f32022-01-12 07:19:28 -080030#include "google/protobuf/descriptor.upbdefs.h"
Joshua Haberman8f0a2cd2022-01-07 17:43:53 -080031#include "python/convert.h"
Joshua Haberman94ece042021-08-18 12:38:26 -070032#include "python/descriptor.h"
Joshua Haberman55e5da52021-12-28 21:42:50 -080033#include "python/message.h"
Joshua Haberman94ece042021-08-18 12:38:26 -070034#include "python/protobuf.h"
Eric Saloedecfd52022-09-15 10:26:14 -070035#include "upb/reflection/def.h"
Joshua Haberman8f0a2cd2022-01-07 17:43:53 -080036#include "upb/util/def_to_proto.h"
Joshua Haberman94ece042021-08-18 12:38:26 -070037
38// -----------------------------------------------------------------------------
39// DescriptorPool
40// -----------------------------------------------------------------------------
41
42typedef struct {
Joshua Haberman1c955f32022-01-12 07:19:28 -080043 PyObject_HEAD;
44 upb_DefPool* symtab;
Joshua Habermancc9b4862021-12-29 11:48:23 -080045 PyObject* db; // The DescriptorDatabase underlying this pool. May be NULL.
Joshua Haberman94ece042021-08-18 12:38:26 -070046} PyUpb_DescriptorPool;
47
Protobuf Team Botb8f35b22023-01-04 13:03:04 -080048PyObject* PyUpb_DescriptorPool_GetDefaultPool(void) {
Joshua Habermane3cd0582021-12-02 15:00:55 -080049 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
Joshua Habermandeddc202021-12-01 15:11:12 -080050 return s->default_pool;
51}
52
Joshua Haberman1c955f32022-01-12 07:19:28 -080053const upb_MessageDef* PyUpb_DescriptorPool_GetFileProtoDef(void) {
Joshua Haberman8f0a2cd2022-01-07 17:43:53 -080054 PyUpb_ModuleState* s = PyUpb_ModuleState_Get();
55 if (!s->c_descriptor_symtab) {
Joshua Haberman1c955f32022-01-12 07:19:28 -080056 s->c_descriptor_symtab = upb_DefPool_New();
Joshua Haberman8f0a2cd2022-01-07 17:43:53 -080057 }
58 return google_protobuf_FileDescriptorProto_getmsgdef(s->c_descriptor_symtab);
59}
60
Joshua Haberman55e5da52021-12-28 21:42:50 -080061static PyObject* PyUpb_DescriptorPool_DoCreateWithCache(
62 PyTypeObject* type, PyObject* db, PyUpb_WeakMap* obj_cache) {
Joshua Habermaneb38a2a2021-12-05 23:54:34 -080063 PyUpb_DescriptorPool* pool = (void*)PyType_GenericAlloc(type, 0);
Joshua Haberman1c955f32022-01-12 07:19:28 -080064 pool->symtab = upb_DefPool_New();
Joshua Haberman94ece042021-08-18 12:38:26 -070065 pool->db = db;
Joshua Habermancdc9fcd2021-08-23 16:57:40 -070066 Py_XINCREF(pool->db);
Joshua Habermaneb38a2a2021-12-05 23:54:34 -080067 PyUpb_WeakMap_Add(obj_cache, pool->symtab, &pool->ob_base);
Joshua Haberman94ece042021-08-18 12:38:26 -070068 return &pool->ob_base;
69}
70
Joshua Habermaneb38a2a2021-12-05 23:54:34 -080071static PyObject* PyUpb_DescriptorPool_DoCreate(PyTypeObject* type,
72 PyObject* db) {
73 return PyUpb_DescriptorPool_DoCreateWithCache(type, db,
74 PyUpb_ObjCache_Instance());
75}
76
Joshua Haberman1c955f32022-01-12 07:19:28 -080077upb_DefPool* PyUpb_DescriptorPool_GetSymtab(PyObject* pool) {
Joshua Habermandeddc202021-12-01 15:11:12 -080078 return ((PyUpb_DescriptorPool*)pool)->symtab;
79}
80
Joshua Habermancdc9fcd2021-08-23 16:57:40 -070081static int PyUpb_DescriptorPool_Traverse(PyUpb_DescriptorPool* self,
82 visitproc visit, void* arg) {
83 Py_VISIT(self->db);
84 return 0;
85}
86
87static int PyUpb_DescriptorPool_Clear(PyUpb_DescriptorPool* self) {
88 Py_CLEAR(self->db);
89 return 0;
90}
91
Joshua Haberman1c955f32022-01-12 07:19:28 -080092PyObject* PyUpb_DescriptorPool_Get(const upb_DefPool* symtab) {
Joshua Habermane3cd0582021-12-02 15:00:55 -080093 PyObject* pool = PyUpb_ObjCache_Get(symtab);
Joshua Habermandeddc202021-12-01 15:11:12 -080094 assert(pool);
95 return pool;
96}
97
Joshua Habermaneec1a452021-08-18 15:23:57 -070098static void PyUpb_DescriptorPool_Dealloc(PyUpb_DescriptorPool* self) {
Joshua Habermancdc9fcd2021-08-23 16:57:40 -070099 PyUpb_DescriptorPool_Clear(self);
Joshua Haberman1c955f32022-01-12 07:19:28 -0800100 upb_DefPool_Free(self->symtab);
Joshua Habermandeddc202021-12-01 15:11:12 -0800101 PyUpb_ObjCache_Delete(self->symtab);
102 PyUpb_Dealloc(self);
Joshua Haberman94ece042021-08-18 12:38:26 -0700103}
104
105/*
106 * DescriptorPool.__new__()
107 *
108 * Implements:
109 * DescriptorPool(descriptor_db=None)
110 */
111static PyObject* PyUpb_DescriptorPool_New(PyTypeObject* type, PyObject* args,
112 PyObject* kwargs) {
113 char* kwlist[] = {"descriptor_db", 0};
114 PyObject* db = NULL;
115
116 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O", kwlist, &db)) {
117 return NULL;
118 }
119
Joshua Haberman55e5da52021-12-28 21:42:50 -0800120 if (db == Py_None) db = NULL;
Joshua Haberman94ece042021-08-18 12:38:26 -0700121 return PyUpb_DescriptorPool_DoCreate(type, db);
122}
123
Joshua Habermana754a372022-01-08 12:19:08 -0800124static PyObject* PyUpb_DescriptorPool_DoAdd(PyObject* _self,
125 PyObject* file_desc);
126
127static bool PyUpb_DescriptorPool_TryLoadFileProto(PyUpb_DescriptorPool* self,
128 PyObject* proto) {
129 if (proto == NULL) {
130 if (PyErr_ExceptionMatches(PyExc_KeyError)) {
131 // Expected error: item was simply not found.
132 PyErr_Clear();
133 return true; // We didn't accomplish our goal, but we didn't error out.
134 }
135 return false;
136 }
137 if (proto == Py_None) return true;
Joshua Habermanffdcc462022-01-12 08:38:36 -0800138 PyObject* ret = PyUpb_DescriptorPool_DoAdd((PyObject*)self, proto);
139 bool ok = ret != NULL;
140 Py_XDECREF(ret);
141 return ok;
Joshua Habermana754a372022-01-08 12:19:08 -0800142}
143
144static bool PyUpb_DescriptorPool_TryLoadSymbol(PyUpb_DescriptorPool* self,
145 PyObject* sym) {
146 if (!self->db) return false;
147 PyObject* file_proto =
148 PyObject_CallMethod(self->db, "FindFileContainingSymbol", "O", sym);
149 bool ret = PyUpb_DescriptorPool_TryLoadFileProto(self, file_proto);
150 Py_XDECREF(file_proto);
151 return ret;
152}
153
Joshua Haberman7a243402022-01-08 13:21:14 -0800154static bool PyUpb_DescriptorPool_TryLoadFilename(PyUpb_DescriptorPool* self,
155 PyObject* filename) {
156 if (!self->db) return false;
157 PyObject* file_proto =
158 PyObject_CallMethod(self->db, "FindFileByName", "O", filename);
159 bool ret = PyUpb_DescriptorPool_TryLoadFileProto(self, file_proto);
160 Py_XDECREF(file_proto);
161 return ret;
162}
163
Joshua Haberman1c955f32022-01-12 07:19:28 -0800164bool PyUpb_DescriptorPool_CheckNoDatabase(PyObject* _self) { return true; }
Joshua Haberman00c106f2022-01-08 22:15:10 -0800165
166static bool PyUpb_DescriptorPool_LoadDependentFiles(
167 PyUpb_DescriptorPool* self, google_protobuf_FileDescriptorProto* proto) {
Joshua Haberman00c106f2022-01-08 22:15:10 -0800168 size_t n;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800169 const upb_StringView* deps =
Joshua Haberman00c106f2022-01-08 22:15:10 -0800170 google_protobuf_FileDescriptorProto_dependency(proto, &n);
171 for (size_t i = 0; i < n; i++) {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800172 const upb_FileDef* dep = upb_DefPool_FindFileByNameWithSize(
173 self->symtab, deps[i].data, deps[i].size);
Joshua Haberman00c106f2022-01-08 22:15:10 -0800174 if (!dep) {
175 PyObject* filename =
176 PyUnicode_FromStringAndSize(deps[i].data, deps[i].size);
177 if (!filename) return false;
178 bool ok = PyUpb_DescriptorPool_TryLoadFilename(self, filename);
179 Py_DECREF(filename);
180 if (!ok) return false;
181 }
Joshua Habermana754a372022-01-08 12:19:08 -0800182 }
183 return true;
184}
185
186static PyObject* PyUpb_DescriptorPool_DoAddSerializedFile(
Joshua Haberman94ece042021-08-18 12:38:26 -0700187 PyObject* _self, PyObject* serialized_pb) {
188 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800189 upb_Arena* arena = upb_Arena_New();
Joshua Habermane3cd0582021-12-02 15:00:55 -0800190 if (!arena) PYUPB_RETURN_OOM;
Joshua Haberman94ece042021-08-18 12:38:26 -0700191 PyObject* result = NULL;
192
Joshua Haberman00c106f2022-01-08 22:15:10 -0800193 char* buf;
194 Py_ssize_t size;
Joshua Haberman94ece042021-08-18 12:38:26 -0700195 if (PyBytes_AsStringAndSize(serialized_pb, &buf, &size) < 0) {
196 goto done;
197 }
198
199 google_protobuf_FileDescriptorProto* proto =
200 google_protobuf_FileDescriptorProto_parse(buf, size, arena);
201 if (!proto) {
202 PyErr_SetString(PyExc_TypeError, "Couldn't parse file content!");
Joshua Habermancdc9fcd2021-08-23 16:57:40 -0700203 goto done;
Joshua Haberman94ece042021-08-18 12:38:26 -0700204 }
205
Joshua Haberman1c955f32022-01-12 07:19:28 -0800206 upb_StringView name = google_protobuf_FileDescriptorProto_name(proto);
207 const upb_FileDef* file =
208 upb_DefPool_FindFileByNameWithSize(self->symtab, name.data, name.size);
Joshua Haberman8f0a2cd2022-01-07 17:43:53 -0800209
210 if (file) {
211 // If the existing file is equal to the new file, then silently ignore the
212 // duplicate add.
213 google_protobuf_FileDescriptorProto* existing =
214 upb_FileDef_ToProto(file, arena);
215 if (!existing) {
216 PyErr_SetNone(PyExc_MemoryError);
217 goto done;
218 }
Joshua Haberman1c955f32022-01-12 07:19:28 -0800219 const upb_MessageDef* m = PyUpb_DescriptorPool_GetFileProtoDef();
Joshua Haberman12fffeb2022-05-17 14:00:49 -0700220 if (upb_Message_IsEqual(proto, existing, m)) {
Eric Saloe779b9d2022-09-27 18:02:21 -0700221 result = PyUpb_FileDescriptor_Get(file);
Joshua Haberman8f0a2cd2022-01-07 17:43:53 -0800222 goto done;
223 }
224 }
225
Joshua Haberman7a243402022-01-08 13:21:14 -0800226 if (self->db) {
Joshua Haberman00c106f2022-01-08 22:15:10 -0800227 if (!PyUpb_DescriptorPool_LoadDependentFiles(self, proto)) goto done;
Joshua Haberman7a243402022-01-08 13:21:14 -0800228 }
229
Joshua Haberman1c955f32022-01-12 07:19:28 -0800230 upb_Status status;
231 upb_Status_Clear(&status);
Joshua Haberman94ece042021-08-18 12:38:26 -0700232
Joshua Haberman1c955f32022-01-12 07:19:28 -0800233 const upb_FileDef* filedef =
234 upb_DefPool_AddFile(self->symtab, proto, &status);
Joshua Haberman94ece042021-08-18 12:38:26 -0700235 if (!filedef) {
236 PyErr_Format(PyExc_TypeError,
237 "Couldn't build proto file into descriptor pool: %s",
Joshua Haberman1c955f32022-01-12 07:19:28 -0800238 upb_Status_ErrorMessage(&status));
Joshua Habermancdc9fcd2021-08-23 16:57:40 -0700239 goto done;
Joshua Haberman94ece042021-08-18 12:38:26 -0700240 }
241
Joshua Habermandeddc202021-12-01 15:11:12 -0800242 result = PyUpb_FileDescriptor_Get(filedef);
Joshua Haberman94ece042021-08-18 12:38:26 -0700243
244done:
Joshua Haberman1c955f32022-01-12 07:19:28 -0800245 upb_Arena_Free(arena);
Joshua Haberman94ece042021-08-18 12:38:26 -0700246 return result;
247}
248
Joshua Habermana754a372022-01-08 12:19:08 -0800249static PyObject* PyUpb_DescriptorPool_DoAdd(PyObject* _self,
250 PyObject* file_desc) {
Joshua Haberman12fffeb2022-05-17 14:00:49 -0700251 if (!PyUpb_Message_Verify(file_desc)) return NULL;
252 const upb_MessageDef* m = PyUpb_Message_GetMsgdef(file_desc);
Joshua Haberman6a94a382022-02-27 17:35:52 -0800253 const char* file_proto_name =
254 PYUPB_DESCRIPTOR_PROTO_PACKAGE ".FileDescriptorProto";
Joshua Haberman1c955f32022-01-12 07:19:28 -0800255 if (strcmp(upb_MessageDef_FullName(m), file_proto_name) != 0) {
Joshua Habermana754a372022-01-08 12:19:08 -0800256 return PyErr_Format(PyExc_TypeError, "Can only add FileDescriptorProto");
257 }
Joshua Habermanffdcc462022-01-12 08:38:36 -0800258 PyObject* subargs = PyTuple_New(0);
259 if (!subargs) return NULL;
Joshua Haberman55e5da52021-12-28 21:42:50 -0800260 PyObject* serialized =
Joshua Haberman12fffeb2022-05-17 14:00:49 -0700261 PyUpb_Message_SerializeToString(file_desc, subargs, NULL);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800262 Py_DECREF(subargs);
263 if (!serialized) return NULL;
Joshua Habermana754a372022-01-08 12:19:08 -0800264 PyObject* ret = PyUpb_DescriptorPool_DoAddSerializedFile(_self, serialized);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800265 Py_DECREF(serialized);
266 return ret;
267}
268
269/*
Joshua Habermana754a372022-01-08 12:19:08 -0800270 * PyUpb_DescriptorPool_AddSerializedFile()
271 *
272 * Implements:
273 * DescriptorPool.AddSerializedFile(self, serialized_file_descriptor)
274 *
275 * Adds the given serialized FileDescriptorProto to the pool.
276 */
277static PyObject* PyUpb_DescriptorPool_AddSerializedFile(
Joshua Haberman1c955f32022-01-12 07:19:28 -0800278 PyObject* _self, PyObject* serialized_pb) {
Joshua Haberman00c106f2022-01-08 22:15:10 -0800279 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
280 if (self->db) {
281 PyErr_SetString(
282 PyExc_ValueError,
283 "Cannot call AddSerializedFile on a DescriptorPool that uses a "
284 "DescriptorDatabase. Add your file to the underlying database.");
285 return false;
286 }
287 return PyUpb_DescriptorPool_DoAddSerializedFile(_self, serialized_pb);
Joshua Habermana754a372022-01-08 12:19:08 -0800288}
289
Joshua Haberman00c106f2022-01-08 22:15:10 -0800290static PyObject* PyUpb_DescriptorPool_Add(PyObject* _self,
Joshua Habermana754a372022-01-08 12:19:08 -0800291 PyObject* file_desc) {
Joshua Haberman00c106f2022-01-08 22:15:10 -0800292 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
293 if (self->db) {
294 PyErr_SetString(
295 PyExc_ValueError,
296 "Cannot call Add on a DescriptorPool that uses a DescriptorDatabase. "
297 "Add your file to the underlying database.");
298 return false;
299 }
300 return PyUpb_DescriptorPool_DoAdd(_self, file_desc);
Joshua Habermana754a372022-01-08 12:19:08 -0800301}
302
303/*
Joshua Haberman55e5da52021-12-28 21:42:50 -0800304 * PyUpb_DescriptorPool_FindFileByName()
305 *
306 * Implements:
307 * DescriptorPool.FindFileByName(self, name)
308 */
309static PyObject* PyUpb_DescriptorPool_FindFileByName(PyObject* _self,
310 PyObject* arg) {
311 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
312
Joshua Habermancfdd0162022-01-16 21:32:49 -0800313 const char* name = PyUpb_VerifyStrData(arg);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800314 if (!name) return NULL;
315
Joshua Haberman1c955f32022-01-12 07:19:28 -0800316 const upb_FileDef* file = upb_DefPool_FindFileByName(self->symtab, name);
Joshua Haberman7a243402022-01-08 13:21:14 -0800317 if (file == NULL && self->db) {
318 if (!PyUpb_DescriptorPool_TryLoadFilename(self, arg)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800319 file = upb_DefPool_FindFileByName(self->symtab, name);
Joshua Haberman7a243402022-01-08 13:21:14 -0800320 }
Joshua Haberman55e5da52021-12-28 21:42:50 -0800321 if (file == NULL) {
322 return PyErr_Format(PyExc_KeyError, "Couldn't find file %.200s", name);
323 }
324
325 return PyUpb_FileDescriptor_Get(file);
326}
327
Joshua Haberman94ece042021-08-18 12:38:26 -0700328/*
329 * PyUpb_DescriptorPool_FindExtensionByName()
330 *
331 * Implements:
332 * DescriptorPool.FindExtensionByName(self, name)
333 */
334static PyObject* PyUpb_DescriptorPool_FindExtensionByName(PyObject* _self,
335 PyObject* arg) {
336 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
337
Joshua Habermancfdd0162022-01-16 21:32:49 -0800338 const char* name = PyUpb_VerifyStrData(arg);
Joshua Habermancdc9fcd2021-08-23 16:57:40 -0700339 if (!name) return NULL;
Joshua Haberman94ece042021-08-18 12:38:26 -0700340
Joshua Haberman1c955f32022-01-12 07:19:28 -0800341 const upb_FieldDef* field =
342 upb_DefPool_FindExtensionByName(self->symtab, name);
Joshua Haberman7a243402022-01-08 13:21:14 -0800343 if (field == NULL && self->db) {
344 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800345 field = upb_DefPool_FindExtensionByName(self->symtab, name);
Joshua Haberman7a243402022-01-08 13:21:14 -0800346 }
Joshua Haberman94ece042021-08-18 12:38:26 -0700347 if (field == NULL) {
348 return PyErr_Format(PyExc_KeyError, "Couldn't find extension %.200s", name);
349 }
350
Joshua Habermandeddc202021-12-01 15:11:12 -0800351 return PyUpb_FieldDescriptor_Get(field);
Joshua Haberman94ece042021-08-18 12:38:26 -0700352}
353
Joshua Haberman55e5da52021-12-28 21:42:50 -0800354/*
355 * PyUpb_DescriptorPool_FindMessageTypeByName()
356 *
357 * Implements:
358 * DescriptorPool.FindMessageTypeByName(self, name)
359 */
360static PyObject* PyUpb_DescriptorPool_FindMessageTypeByName(PyObject* _self,
361 PyObject* arg) {
362 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
363
Joshua Habermancfdd0162022-01-16 21:32:49 -0800364 const char* name = PyUpb_VerifyStrData(arg);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800365 if (!name) return NULL;
366
Joshua Haberman1c955f32022-01-12 07:19:28 -0800367 const upb_MessageDef* m = upb_DefPool_FindMessageByName(self->symtab, name);
Joshua Habermana754a372022-01-08 12:19:08 -0800368 if (m == NULL && self->db) {
369 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800370 m = upb_DefPool_FindMessageByName(self->symtab, name);
Joshua Habermana754a372022-01-08 12:19:08 -0800371 }
Joshua Haberman55e5da52021-12-28 21:42:50 -0800372 if (m == NULL) {
373 return PyErr_Format(PyExc_KeyError, "Couldn't find message %.200s", name);
374 }
375
376 return PyUpb_Descriptor_Get(m);
377}
378
Joshua Habermane4c88ad2021-12-29 14:41:44 -0800379// Splits a dotted symbol like foo.bar.baz on the last dot. Returns the portion
380// after the last dot (baz) and updates `*parent_size` to the length of the
381// parent (foo.bar). Returns NULL if no dots were present.
382static const char* PyUpb_DescriptorPool_SplitSymbolName(const char* sym,
383 size_t* parent_size) {
384 const char* last_dot = strrchr(sym, '.');
385 if (!last_dot) return NULL;
386 *parent_size = last_dot - sym;
387 return last_dot + 1;
388}
389
Joshua Haberman55e5da52021-12-28 21:42:50 -0800390/*
391 * PyUpb_DescriptorPool_FindFieldByName()
392 *
393 * Implements:
394 * DescriptorPool.FindFieldByName(self, name)
395 */
396static PyObject* PyUpb_DescriptorPool_FindFieldByName(PyObject* _self,
397 PyObject* arg) {
398 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
399
Joshua Habermancfdd0162022-01-16 21:32:49 -0800400 const char* name = PyUpb_VerifyStrData(arg);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800401 if (!name) return NULL;
402
Joshua Haberman89ff28c2022-01-07 17:18:55 -0800403 size_t parent_size;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800404 const char* child = PyUpb_DescriptorPool_SplitSymbolName(name, &parent_size);
405 const upb_FieldDef* f = NULL;
Joshua Haberman89ff28c2022-01-07 17:18:55 -0800406 if (child) {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800407 const upb_MessageDef* parent =
408 upb_DefPool_FindMessageByNameWithSize(self->symtab, name, parent_size);
Joshua Haberman3eeea892022-01-16 18:02:06 -0800409 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 }
Joshua Haberman89ff28c2022-01-07 17:18:55 -0800414 if (parent) {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800415 f = upb_MessageDef_FindFieldByName(parent, child);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800416 }
417 }
418
419 if (!f) {
420 return PyErr_Format(PyExc_KeyError, "Couldn't find message %.200s", name);
421 }
422
423 return PyUpb_FieldDescriptor_Get(f);
424}
425
426/*
427 * PyUpb_DescriptorPool_FindEnumTypeByName()
428 *
429 * Implements:
430 * DescriptorPool.FindEnumTypeByName(self, name)
431 */
432static PyObject* PyUpb_DescriptorPool_FindEnumTypeByName(PyObject* _self,
433 PyObject* arg) {
434 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
435
Joshua Habermancfdd0162022-01-16 21:32:49 -0800436 const char* name = PyUpb_VerifyStrData(arg);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800437 if (!name) return NULL;
438
Joshua Haberman1c955f32022-01-12 07:19:28 -0800439 const upb_EnumDef* e = upb_DefPool_FindEnumByName(self->symtab, name);
Joshua Haberman7a243402022-01-08 13:21:14 -0800440 if (e == NULL && self->db) {
441 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800442 e = upb_DefPool_FindEnumByName(self->symtab, name);
Joshua Haberman7a243402022-01-08 13:21:14 -0800443 }
Joshua Haberman55e5da52021-12-28 21:42:50 -0800444 if (e == NULL) {
445 return PyErr_Format(PyExc_KeyError, "Couldn't find enum %.200s", name);
446 }
447
448 return PyUpb_EnumDescriptor_Get(e);
449}
450
451/*
452 * PyUpb_DescriptorPool_FindOneofByName()
453 *
454 * Implements:
455 * DescriptorPool.FindOneofByName(self, name)
456 */
457static PyObject* PyUpb_DescriptorPool_FindOneofByName(PyObject* _self,
458 PyObject* arg) {
459 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
460
Joshua Habermancfdd0162022-01-16 21:32:49 -0800461 const char* name = PyUpb_VerifyStrData(arg);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800462 if (!name) return NULL;
463
Joshua Habermane4c88ad2021-12-29 14:41:44 -0800464 size_t parent_size;
465 const char* child = PyUpb_DescriptorPool_SplitSymbolName(name, &parent_size);
466
467 if (child) {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800468 const upb_MessageDef* parent =
469 upb_DefPool_FindMessageByNameWithSize(self->symtab, name, parent_size);
Joshua Haberman7a243402022-01-08 13:21:14 -0800470 if (parent == NULL && self->db) {
471 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800472 parent = upb_DefPool_FindMessageByNameWithSize(self->symtab, name,
473 parent_size);
Joshua Haberman7a243402022-01-08 13:21:14 -0800474 }
Joshua Haberman55e5da52021-12-28 21:42:50 -0800475 if (parent) {
Joshua Haberman1c955f32022-01-12 07:19:28 -0800476 const upb_OneofDef* o = upb_MessageDef_FindOneofByName(parent, child);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800477 return PyUpb_OneofDescriptor_Get(o);
478 }
479 }
480
Joshua Haberman7a243402022-01-08 13:21:14 -0800481 return PyErr_Format(PyExc_KeyError, "Couldn't find oneof %.200s", name);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800482}
483
484static PyObject* PyUpb_DescriptorPool_FindServiceByName(PyObject* _self,
485 PyObject* arg) {
486 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
487
Joshua Habermancfdd0162022-01-16 21:32:49 -0800488 const char* name = PyUpb_VerifyStrData(arg);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800489 if (!name) return NULL;
490
Joshua Haberman1c955f32022-01-12 07:19:28 -0800491 const upb_ServiceDef* s = upb_DefPool_FindServiceByName(self->symtab, name);
Joshua Haberman392531c2022-01-16 18:25:29 -0800492 if (s == NULL && self->db) {
493 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
494 s = upb_DefPool_FindServiceByName(self->symtab, name);
495 }
Joshua Haberman55e5da52021-12-28 21:42:50 -0800496 if (s == NULL) {
Joshua Haberman392531c2022-01-16 18:25:29 -0800497 return PyErr_Format(PyExc_KeyError, "Couldn't find service %.200s", name);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800498 }
499
500 return PyUpb_ServiceDescriptor_Get(s);
501}
502
Joshua Haberman392531c2022-01-16 18:25:29 -0800503static PyObject* PyUpb_DescriptorPool_FindMethodByName(PyObject* _self,
504 PyObject* arg) {
505 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
506
Joshua Habermancfdd0162022-01-16 21:32:49 -0800507 const char* name = PyUpb_VerifyStrData(arg);
Joshua Haberman392531c2022-01-16 18:25:29 -0800508 if (!name) return NULL;
509 size_t parent_size;
510 const char* child = PyUpb_DescriptorPool_SplitSymbolName(name, &parent_size);
511
512 if (!child) goto err;
513 const upb_ServiceDef* parent =
514 upb_DefPool_FindServiceByNameWithSize(self->symtab, name, parent_size);
515 if (parent == NULL && self->db) {
516 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
Joshua Haberman75b62912022-01-16 19:03:45 -0800517 parent =
518 upb_DefPool_FindServiceByNameWithSize(self->symtab, name, parent_size);
Joshua Haberman392531c2022-01-16 18:25:29 -0800519 }
520 if (!parent) goto err;
521 const upb_MethodDef* m = upb_ServiceDef_FindMethodByName(parent, child);
522 if (!m) goto err;
523 return PyUpb_MethodDescriptor_Get(m);
524
525err:
526 return PyErr_Format(PyExc_KeyError, "Couldn't find method %.200s", name);
527}
528
Joshua Haberman55e5da52021-12-28 21:42:50 -0800529static PyObject* PyUpb_DescriptorPool_FindFileContainingSymbol(PyObject* _self,
530 PyObject* arg) {
531 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
532
Joshua Habermancfdd0162022-01-16 21:32:49 -0800533 const char* name = PyUpb_VerifyStrData(arg);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800534 if (!name) return NULL;
535
Joshua Haberman392531c2022-01-16 18:25:29 -0800536 const upb_FileDef* f =
537 upb_DefPool_FindFileContainingSymbol(self->symtab, name);
Joshua Haberman7a243402022-01-08 13:21:14 -0800538 if (f == NULL && self->db) {
539 if (!PyUpb_DescriptorPool_TryLoadSymbol(self, arg)) return NULL;
Joshua Haberman392531c2022-01-16 18:25:29 -0800540 f = upb_DefPool_FindFileContainingSymbol(self->symtab, name);
Joshua Haberman7a243402022-01-08 13:21:14 -0800541 }
Joshua Haberman55e5da52021-12-28 21:42:50 -0800542 if (f == NULL) {
543 return PyErr_Format(PyExc_KeyError, "Couldn't find symbol %.200s", name);
544 }
545
546 return PyUpb_FileDescriptor_Get(f);
547}
548
549static PyObject* PyUpb_DescriptorPool_FindExtensionByNumber(PyObject* _self,
550 PyObject* args) {
551 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
552 PyObject* message_descriptor;
553 int number;
554 if (!PyArg_ParseTuple(args, "Oi", &message_descriptor, &number)) {
555 return NULL;
556 }
557
Joshua Haberman1c955f32022-01-12 07:19:28 -0800558 const upb_FieldDef* f = upb_DefPool_FindExtensionByNumber(
Joshua Haberman55e5da52021-12-28 21:42:50 -0800559 self->symtab, PyUpb_Descriptor_GetDef(message_descriptor), number);
560 if (f == NULL) {
561 return PyErr_Format(PyExc_KeyError, "Couldn't find Extension %d", number);
562 }
563
564 return PyUpb_FieldDescriptor_Get(f);
565}
566
567static PyObject* PyUpb_DescriptorPool_FindAllExtensions(PyObject* _self,
568 PyObject* msg_desc) {
569 PyUpb_DescriptorPool* self = (PyUpb_DescriptorPool*)_self;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800570 const upb_MessageDef* m = PyUpb_Descriptor_GetDef(msg_desc);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800571 size_t n;
Joshua Haberman1c955f32022-01-12 07:19:28 -0800572 const upb_FieldDef** ext = upb_DefPool_GetAllExtensions(self->symtab, m, &n);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800573 PyObject* ret = PyList_New(n);
Joshua Haberman92e06d02022-04-25 09:52:40 -0700574 if (!ret) goto done;
Joshua Haberman55e5da52021-12-28 21:42:50 -0800575 for (size_t i = 0; i < n; i++) {
576 PyObject* field = PyUpb_FieldDescriptor_Get(ext[i]);
Joshua Haberman92e06d02022-04-25 09:52:40 -0700577 if (!field) {
578 Py_DECREF(ret);
579 ret = NULL;
580 goto done;
581 }
Joshua Haberman55e5da52021-12-28 21:42:50 -0800582 PyList_SetItem(ret, i, field);
583 }
Joshua Haberman92e06d02022-04-25 09:52:40 -0700584done:
585 free(ext);
Joshua Haberman55e5da52021-12-28 21:42:50 -0800586 return ret;
587}
588
Joshua Haberman94ece042021-08-18 12:38:26 -0700589static PyMethodDef PyUpb_DescriptorPool_Methods[] = {
Joshua Haberman55e5da52021-12-28 21:42:50 -0800590 {"Add", PyUpb_DescriptorPool_Add, METH_O,
591 "Adds the FileDescriptorProto and its types to this pool."},
Joshua Haberman94ece042021-08-18 12:38:26 -0700592 {"AddSerializedFile", PyUpb_DescriptorPool_AddSerializedFile, METH_O,
593 "Adds a serialized FileDescriptorProto to this pool."},
Joshua Haberman55e5da52021-12-28 21:42:50 -0800594 {"FindFileByName", PyUpb_DescriptorPool_FindFileByName, METH_O,
595 "Searches for a file descriptor by its .proto name."},
596 {"FindMessageTypeByName", PyUpb_DescriptorPool_FindMessageTypeByName,
597 METH_O, "Searches for a message descriptor by full name."},
598 {"FindFieldByName", PyUpb_DescriptorPool_FindFieldByName, METH_O,
599 "Searches for a field descriptor by full name."},
Joshua Haberman94ece042021-08-18 12:38:26 -0700600 {"FindExtensionByName", PyUpb_DescriptorPool_FindExtensionByName, METH_O,
601 "Searches for extension descriptor by full name."},
Joshua Haberman55e5da52021-12-28 21:42:50 -0800602 {"FindEnumTypeByName", PyUpb_DescriptorPool_FindEnumTypeByName, METH_O,
603 "Searches for enum type descriptor by full name."},
604 {"FindOneofByName", PyUpb_DescriptorPool_FindOneofByName, METH_O,
605 "Searches for oneof descriptor by full name."},
606 {"FindServiceByName", PyUpb_DescriptorPool_FindServiceByName, METH_O,
607 "Searches for service descriptor by full name."},
Joshua Haberman392531c2022-01-16 18:25:29 -0800608 {"FindMethodByName", PyUpb_DescriptorPool_FindMethodByName, METH_O,
609 "Searches for method descriptor by full name."},
Joshua Haberman55e5da52021-12-28 21:42:50 -0800610 {"FindFileContainingSymbol", PyUpb_DescriptorPool_FindFileContainingSymbol,
611 METH_O, "Gets the FileDescriptor containing the specified symbol."},
612 {"FindExtensionByNumber", PyUpb_DescriptorPool_FindExtensionByNumber,
613 METH_VARARGS, "Gets the extension descriptor for the given number."},
614 {"FindAllExtensions", PyUpb_DescriptorPool_FindAllExtensions, METH_O,
615 "Gets all known extensions of the given message descriptor."},
Joshua Haberman94ece042021-08-18 12:38:26 -0700616 {NULL}};
617
618static PyType_Slot PyUpb_DescriptorPool_Slots[] = {
Joshua Habermancdc9fcd2021-08-23 16:57:40 -0700619 {Py_tp_clear, PyUpb_DescriptorPool_Clear},
Joshua Haberman94ece042021-08-18 12:38:26 -0700620 {Py_tp_dealloc, PyUpb_DescriptorPool_Dealloc},
621 {Py_tp_methods, PyUpb_DescriptorPool_Methods},
Joshua Habermancdc9fcd2021-08-23 16:57:40 -0700622 {Py_tp_new, PyUpb_DescriptorPool_New},
623 {Py_tp_traverse, PyUpb_DescriptorPool_Traverse},
Joshua Haberman94ece042021-08-18 12:38:26 -0700624 {0, NULL}};
625
626static PyType_Spec PyUpb_DescriptorPool_Spec = {
Joshua Habermancdc9fcd2021-08-23 16:57:40 -0700627 PYUPB_MODULE_NAME ".DescriptorPool",
628 sizeof(PyUpb_DescriptorPool),
629 0, // tp_itemsize
630 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
Joshua Haberman94ece042021-08-18 12:38:26 -0700631 PyUpb_DescriptorPool_Slots,
632};
633
634// -----------------------------------------------------------------------------
635// Top Level
636// -----------------------------------------------------------------------------
637
638bool PyUpb_InitDescriptorPool(PyObject* m) {
Joshua Haberman55e5da52021-12-28 21:42:50 -0800639 PyUpb_ModuleState* state = PyUpb_ModuleState_GetFromModule(m);
Joshua Haberman94ece042021-08-18 12:38:26 -0700640 PyTypeObject* descriptor_pool_type =
Joshua Haberman55e5da52021-12-28 21:42:50 -0800641 PyUpb_AddClass(m, &PyUpb_DescriptorPool_Spec);
Joshua Haberman94ece042021-08-18 12:38:26 -0700642
643 if (!descriptor_pool_type) return false;
644
Joshua Habermaneb38a2a2021-12-05 23:54:34 -0800645 state->default_pool = PyUpb_DescriptorPool_DoCreateWithCache(
646 descriptor_pool_type, NULL, state->obj_cache);
647 return state->default_pool &&
648 PyModule_AddObject(m, "default_pool", state->default_pool) == 0;
Joshua Haberman94ece042021-08-18 12:38:26 -0700649}