diff --git a/swh/model/from_disk.py b/swh/model/from_disk.py --- a/swh/model/from_disk.py +++ b/swh/model/from_disk.py @@ -549,6 +549,13 @@ key1, key2 = key.rsplit(b"/", 1) del self[key1][key2] + def __contains__(self, key): + if b"/" not in key: + return super().__contains__(key) + else: + key1, key2 = key.split(b"/", 1) + return super().__contains__(key1) and self[key1].__contains__(key2) + def __repr__(self): return "Directory(id=%s, entries=[%s])" % ( hash_to_hex(self.hash), diff --git a/swh/model/tests/test_from_disk.py b/swh/model/tests/test_from_disk.py --- a/swh/model/tests/test_from_disk.py +++ b/swh/model/tests/test_from_disk.py @@ -978,3 +978,19 @@ del d["foo"] with self.assertRaisesRegex(ValueError, "bytes Directory entry"): del d[42] + + def test_directory_contains(self): + d = Directory() + d[b"a"] = Directory() + d[b"a/b"] = Directory() + d[b"a/b/c"] = Directory() + d[b"a/b/c/d"] = Content() + + self.assertIn(b"a", d) + self.assertIn(b"a/b", d) + self.assertIn(b"a/b/c", d) + self.assertIn(b"a/b/c/d", d) + + self.assertNotIn(b"b", d) + self.assertNotIn(b"b/c", d) + self.assertNotIn(b"b/c/d", d)