Merge pull request #31 from blapid/master

Search functionality for containers

construct/lib/container.py

@@ -70,6 +70,35 @@ def copy(self):
         inst.update(self.iteritems())
         return inst
 
+    def _search(self, name, search_all):
+        items = []
+        for key in self.keys():
+            try:
+                if key == name:
+                    if search_all:
+                        items.append(self[key])
+                    else:
+                        return self[key]
+                if type(self[key]) == Container or type(self[key]) == ListContainer:
+                    ret = self[key]._search(name, search_all)
+                    if ret is not None:
+                        if search_all:
+                            items.extend(ret)
+                        else:
+                            return ret
+            except:
+                pass
+        if search_all:
+            return items
+        else:    
+            return None
+
+    def search(self, name):
+        return self._search(name, False)
+
+    def search_all(self, name):
+        return self._search(name, True)
+
     __update__ = update
     __copy__ = copy
 
@@ -157,6 +186,29 @@ def __pretty_str__(self, nesting = 1, indentation = "    "):
         lines.append(indentation * (nesting - 1))
         lines.append("]")
         return "".join(lines)
+
+    def _search(self, name, search_all):
+        items = []
+        for item in self:
+            try:
+                ret = item._search(name, search_all)
+            except:
+                continue
+            if ret is not None:
+                if search_all:
+                    items.extend(ret)
+                else:
+                    return ret
+        if search_all:
+            return items
+        else:
+            return None
+
+    def search(self, name):
+        return self._search(name, False)
+
+    def search_all(self, name):
+        return self._search(name, True)
 
 
 class LazyContainer(object):

tests/test_search.py

@@ -0,0 +1,77 @@
+import unittest
+from construct import Struct, UBInt8, Bytes, Container, Switch, GreedyRange
+import six
+
+struct = Struct("a",
+    UBInt8("aa"),
+    Struct("ab",
+        UBInt8("aba"),
+        UBInt8("abb"),
+        Struct("abc",
+            UBInt8("abca"),
+            Switch("abcb", lambda ctx: ctx.abca, {
+                    1 : Struct("abcb1",
+                            UBInt8("abcb1a")
+                        ),
+                    2 : Struct("abcb2",
+                            UBInt8("abcb2a")
+                        ),
+                    3 : Struct("abcb3",
+                            UBInt8("abcb3a")
+                        ),
+                    4 : Struct("abcb4",
+                            UBInt8("abcb4a")
+                        ),
+                }
+            ),
+        ),
+    ),
+    UBInt8("ac"),
+    GreedyRange(Struct('ad', UBInt8("ada")))
+)
+
+class TestSearch(unittest.TestCase):
+    def test_search_sanity(self):
+        obj1 = struct.parse(six.b("\x11\x21\x22\x02\x02\x13\x51\x52"))
+
+        self.assertEqual(obj1.search("bb"), None)
+        self.assertNotEqual(obj1.search("abcb"), None)
+        self.assertNotEqual(obj1.search("ad"), None)
+        self.assertEqual(obj1.search("aa"), 0x11)
+        self.assertEqual(obj1.search("aba"), 0x21)
+        self.assertEqual(obj1.search("abb"), 0x22)
+        self.assertEqual(obj1.search('ac'), 0x13)
+
+    def test_search_functionality(self):
+        obj1 = struct.parse(six.b("\x11\x21\x22\x02\x02\x13\x51\x52"))
+        obj2 = struct.parse(six.b("\x11\x21\x22\x03\x03\x13\x51\x52"))
+
+        self.assertEqual(obj1.search('abcb1a'), None)
+        self.assertEqual(obj1.search('abcb3a'), None)
+        self.assertEqual(obj1.search('abcb4a'), None)
+        self.assertEqual(obj1.search('abcb2a'), 0x02)
+
+        self.assertEqual(obj2.search('abcb1a'), None)
+        self.assertEqual(obj2.search('abcb2a'), None)
+        self.assertEqual(obj2.search('abcb4a'), None)
+        self.assertEqual(obj2.search('abcb3a'), 0x03)
+        # Return only the first one
+        self.assertEqual(obj1.search("ada"), 0x51)
+
+    def test_search_all_sanity(self):
+        obj1 = struct.parse(six.b("\x11\x21\x22\x02\x02\x13\x51\x52"))
+
+        self.assertEqual(obj1.search_all("bb"), [])
+        self.assertNotEqual(obj1.search_all("ad"), None)
+        self.assertEqual(obj1.search_all("aa"), [0x11])
+        self.assertEqual(obj1.search_all("aba"), [0x21])
+        self.assertEqual(obj1.search_all("abb"), [0x22])
+        self.assertEqual(obj1.search_all('ac'), [0x13])
+
+    def test_search_all_functionality(self):
+        obj1 = struct.parse(six.b("\x11\x21\x22\x02\x02\x13\x51\x52"))
+        # Return all of them
+        self.assertEqual(obj1.search_all("ada"), [0x51,0x52])
+
+if __name__ == "__main__":
+    unittest.main()