From 25e7a9d3b80207200ce2d9b1ea7abc23e7ae01b2 Mon Sep 17 00:00:00 2001 From: Omar Sandoval Date: Fri, 28 Jun 2019 16:02:52 -0700 Subject: [PATCH] libdrgn/python: implement Program.__contains__ --- libdrgn/drgn.h | 5 +++-- libdrgn/program.c | 4 ++-- libdrgn/python/program.c | 48 +++++++++++++++++++++++++++++++++++++++- tests/test_program.py | 4 ++++ 4 files changed, 56 insertions(+), 5 deletions(-) diff --git a/libdrgn/drgn.h b/libdrgn/drgn.h index db8ec9a2..b6cf2c73 100644 --- a/libdrgn/drgn.h +++ b/libdrgn/drgn.h @@ -1230,8 +1230,9 @@ struct drgn_error *drgn_program_find_type(struct drgn_program *prog, * matched with @ref drgn_filename_matches(). If multiple definitions match, one * is returned arbitrarily. * @param[in] flags Flags indicating what kind of object to look for. - * @param[out] ret Returned object. It must have already been initialized with - * @ref drgn_object_init(). + * @param[out] ret Returned object. This can be @c NULL to check for the + * object's existence without returning it. If not @c NULL, this must have + * already been initialized with @ref drgn_object_init(). * @return @c NULL on success, non-@c NULL on error. */ struct drgn_error *drgn_program_find_object(struct drgn_program *prog, diff --git a/libdrgn/program.c b/libdrgn/program.c index 047a0462..10fe9270 100644 --- a/libdrgn/program.c +++ b/libdrgn/program.c @@ -1718,14 +1718,14 @@ drgn_program_find_object(struct drgn_program *prog, const char *name, struct drgn_symbol sym; struct drgn_qualified_type qualified_type; - if (ret->prog != prog) { + if (ret && ret->prog != prog) { return drgn_error_create(DRGN_ERROR_INVALID_ARGUMENT, "object is from wrong program"); } err = drgn_symbol_index_find(&prog->sindex, name, filename, flags, &sym); - if (err) + if (err || !ret) return err; qualified_type.type = sym.type; qualified_type.qualifiers = sym.qualifiers; diff --git a/libdrgn/python/program.c b/libdrgn/python/program.c index 56996f97..10d09e0d 100644 --- a/libdrgn/python/program.c +++ b/libdrgn/python/program.c @@ -755,6 +755,38 @@ static DrgnObject *Program_subscript(Program *self, PyObject *key) return ret; } +static int Program_contains(Program *self, PyObject *key) +{ + struct drgn_error *err; + const char *name; + bool clear; + + if (!PyUnicode_Check(key)) { + PyErr_SetObject(PyExc_KeyError, key); + return 0; + } + + name = PyUnicode_AsUTF8(key); + if (!name) + return -1; + + clear = set_drgn_in_python(); + err = drgn_program_find_object(&self->prog, name, NULL, + DRGN_FIND_OBJECT_ANY, NULL); + if (clear) + clear_drgn_in_python(); + if (err) { + if (err->code == DRGN_ERROR_LOOKUP) { + drgn_error_destroy(err); + return 0; + } else { + set_drgn_error(err); + return -1; + } + } + return 1; +} + static PyObject *Program_get_flags(Program *self, void *arg) { return PyObject_CallFunction(ProgramFlags_class, "k", @@ -817,6 +849,20 @@ static PyMappingMethods Program_as_mapping = { (binaryfunc)Program_subscript, /* mp_subscript */ }; + +static PySequenceMethods Program_as_sequence = { + NULL, /* sq_length */ + NULL, /* sq_concat */ + NULL, /* sq_repeat */ + NULL, /* sq_item */ + NULL, /* sq_slice */ + NULL, /* sq_ass_item */ + NULL, /* sq_ass_slice */ + (objobjproc)Program_contains, /* sq_contains */ + NULL, /* sq_inplace_concat */ + NULL, /* sq_inplace_repeat */ +}; + PyTypeObject Program_type = { PyVarObject_HEAD_INIT(NULL, 0) "_drgn.Program", /* tp_name */ @@ -829,7 +875,7 @@ PyTypeObject Program_type = { NULL, /* tp_as_async */ NULL, /* tp_repr */ NULL, /* tp_as_number */ - NULL, /* tp_as_sequence */ + &Program_as_sequence, /* tp_as_sequence */ &Program_as_mapping, /* tp_as_mapping */ NULL, /* tp_hash */ NULL, /* tp_call */ diff --git a/tests/test_program.py b/tests/test_program.py index de3dd42d..aba4da5b 100644 --- a/tests/test_program.py +++ b/tests/test_program.py @@ -541,6 +541,7 @@ class TestSymbols(unittest.TestCase): self.assertRaises(LookupError, prog._symbol, 'foo', FindObjectFlags.ANY) prog.add_symbol_finder(lambda name, flags, filename: None) self.assertRaises(LookupError, prog._symbol, 'foo', FindObjectFlags.ANY) + self.assertFalse('foo' in prog) def test_constant(self): sym = Symbol(int_type('int', 4, True), value=4096) @@ -548,6 +549,7 @@ class TestSymbols(unittest.TestCase): self.assertEqual(prog._symbol('PAGE_SIZE', FindObjectFlags.CONSTANT), sym) self.assertEqual(prog._symbol('PAGE_SIZE', FindObjectFlags.ANY), sym) + self.assertTrue('PAGE_SIZE' in prog) def test_function(self): sym = Symbol(function_type(void_type(), (), False), address=0xffff0000, @@ -555,6 +557,7 @@ class TestSymbols(unittest.TestCase): prog = mock_program(symbols=[('func', sym)]) self.assertEqual(prog._symbol('func', FindObjectFlags.FUNCTION), sym) self.assertEqual(prog._symbol('func', FindObjectFlags.ANY), sym) + self.assertTrue('func' in prog) def test_variable(self): sym = Symbol(int_type('int', 4, True), address=0xffff0000, @@ -562,6 +565,7 @@ class TestSymbols(unittest.TestCase): prog = mock_program(symbols=[('counter', sym)]) self.assertEqual(prog._symbol('counter', FindObjectFlags.VARIABLE), sym) self.assertEqual(prog._symbol('counter', FindObjectFlags.ANY), sym) + self.assertTrue('counter' in prog) def test_wrong_kind(self): prog = mock_program()