libdrgn/python: implement Program.__contains__

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,

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;

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 */

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()