edges() and iteredges() now work for all applicable dawgs; tests added for all new edges methods

Author: EliFinkelshteyn

dawg_python/dawgs.py

@@ -332,6 +332,24 @@ def items(self, prefix=""):
 
         return res
 
+    def iteritems(self, prefix=""):
+        index = self.dct.ROOT
+        if not isinstance(prefix, bytes):
+            prefix = prefix.encode('utf8')
+        if prefix:
+            index = self.dct.follow_bytes(prefix, index)
+            if not index:
+                return
+
+        completer = wrapper.Completer(self.dct, self.guide)
+        completer.start(index, prefix)
+
+        while completer.next():
+            key, value = completer.key.split(self._payload_separator)
+            # bytes() cast is a python 2.6 fix
+            item = (key.decode('utf8'), a2b_base64(bytes(value)))
+            yield item
+
     def edges(self, prefix=""):
         index = self.dct.ROOT
         if not isinstance(prefix, bytes):
@@ -346,13 +364,15 @@ def edges(self, prefix=""):
         if not edge_follower.start(index, prefix):
             return res
 
-        res.append(edge_follower.decoded_key)
+        vals = self.b_get_value(edge_follower.decoded_key) or [False]
+        res.extend([(edge_follower.decoded_key, val) for val in vals])
         while edge_follower.next():
-            res.append(edge_follower.decoded_key)
+            vals = self.b_get_value(edge_follower.decoded_key) or [False]
+            res.extend([(edge_follower.decoded_key, val) for val in vals])
 
         return res
 
-    def iteritems(self, prefix=""):
+    def iteredges(self, prefix=""):
         index = self.dct.ROOT
         if not isinstance(prefix, bytes):
             prefix = prefix.encode('utf8')
@@ -361,14 +381,17 @@ def iteritems(self, prefix=""):
             if not index:
                 return
 
-        completer = wrapper.Completer(self.dct, self.guide)
-        completer.start(index, prefix)
-
-        while completer.next():
-            key, value = completer.key.split(self._payload_separator)
-            item = (key.decode('utf8'), a2b_base64(bytes(value))) # bytes() cast is a python 2.6 fix
-            yield item
+        edge_follower = wrapper.EdgeFollower(self.dct, self.guide)
+        if not edge_follower.start(index, prefix):
+            return
 
+        vals = self.b_get_value(edge_follower.decoded_key) or [False]
+        for val in vals:
+            yield (edge_follower.decoded_key, val or False)
+        while edge_follower.next():
+            vals = self.b_get_value(edge_follower.decoded_key) or [False]
+            for val in vals:
+                yield (edge_follower.decoded_key, val or False)
 
     def _has_value(self, index):
         return self.dct.follow_bytes(PAYLOAD_SEPARATOR, index)

tests/test_payload_dawg.py

@@ -58,6 +58,18 @@ def test_iterkeys(self):
         d = self.dawg()
         assert list(d.iterkeys()) == d.keys()
 
+    def test_edges(self):
+        d = self.dawg()
+        assert d.edges('foob') == [('fooba', False)]
+        assert d.edges('fooba') == [('foobar', b'data4')]
+        assert d.edges('fo') == [('foo', b'data1'), ('foo', b'data3')]
+
+    def test_iteredges(self):
+        d = self.dawg()
+        assert list(d.iteredges('foob')) == [('fooba', False)]
+        assert list(d.iteredges('fooba')) == [('foobar', b'data4')]
+        assert list(d.iteredges('fo')) == [('foo', b'data1'), ('foo', b'data3')]
+
     def test_key_completion(self):
         d = self.dawg()
         assert d.keys('fo') == ['foo', 'foo', 'foobar']
@@ -75,6 +87,8 @@ def test_iteritems(self):
     def test_items_completion(self):
         d = self.dawg()
         assert d.items('foob') == [('foobar', b'data4')]
+        assert d.items('foo') == [('foo', b'data1'), ('foo', b'data3'),
+                                  ('foobar', b'data4')]
 
     def test_prefixes(self):
         d = self.dawg()