Support syntax: "Name.Foo in dictionary"

2 files changed, 31 insertions, 7 deletions

diff --git a/src/qpdf/object.cpp b/src/qpdf/object.cpp
index 16778cb..4f99f4c 100644
--- a/src/qpdf/object.cpp
+++ b/src/qpdf/object.cpp
@@ -273,6 +273,15 @@ bool operator==(const QPDFObjectHandle& self, const QPDFObjectHandle& other)
 }
 
 
+bool object_has_key(QPDFObjectHandle& h, std::string const& key)
+{
+    if (!h.isDictionary() && !h.isStream())
+        throw py::value_error("object is not a dictionary or a stream");
+    QPDFObjectHandle dict = h.isStream() ? h.getDict() : h;
+    return dict.hasKey(key);
+}
+
+
 QPDFObjectHandle object_get_key(QPDFObjectHandle& h, std::string const& key)
 {
     if (!h.isDictionary() && !h.isStream())
@@ -517,14 +526,15 @@ void init_object(py::module& m)
         )
         .def("keys", &QPDFObjectHandle::getKeys)
         .def("__contains__",
+            [](QPDFObjectHandle &h, QPDFObjectHandle &key) {
+                if (!key.isName())
+                    throw py::type_error("Dictionaries can only contain Names");
+                return object_has_key(h, key.getName());
+            }
+        )
+        .def("__contains__",
             [](QPDFObjectHandle &h, std::string const& key) {
-                if (h.isDictionary()) {
-                    return h.hasKey(key);
-                }
-                if (h.isStream()) {
-                    return h.getDict().hasKey(key);
-                }
-                throw py::value_error("__contains__ not defined for object type");
+                return object_has_key(h, key);
             }
         )
         .def("as_list", &QPDFObjectHandle::getArrayAsVector)
diff --git a/tests/test_object.py b/tests/test_object.py
index 601fa07..67d21af 100644
--- a/tests/test_object.py
+++ b/tests/test_object.py
@@ -160,6 +160,20 @@ def test_name_equality():
     assert Name.Foo == Name('/Foo')
 
 
+def test_dictionary_contains():
+    d = Dictionary({
+        '/Monty': 'Python',
+        '/Flying': 'Circus'
+    })
+    assert Name.Flying in d
+    assert Name('/Monty') in d
+    assert Name.Brian not in d
+
+
+def test_unslashed_name():
+    with pytest.raises(ValueError, match='must begin with'):
+        Name('Monty') not in d
+
 class TestHashViolation:
 
     def check(self, a, b):