edges() working with unicode; added working iteredges(); added basic tests

Author: EliFinkelshteyn

dawg_python/dawgs.py

@@ -153,14 +153,27 @@ def edges(self, prefix=""):
         if not completer.start_edges(index, b_prefix):
             return res
 
-        key = completer.key.decode('utf8')
-        res.append(key)
+        res.append(completer.decoded_key)
         while completer.next_edge():
-            key = completer.key.decode('utf8')
-            res.append(key)
+            res.append(completer.decoded_key)
 
         return res
 
+    def iteredges(self, prefix=""):
+        b_prefix = prefix.encode('utf8')
+
+        index = self.dct.follow_bytes(b_prefix, self.dct.ROOT)
+        if index is None:
+            return
+
+        completer = wrapper.Completer(self.dct, self.guide)
+        if not completer.start_edges(index, b_prefix):
+            return
+
+        yield completer.decoded_key
+        while completer.next_edge():
+            yield completer.decoded_key
+
     def iterkeys(self, prefix=""):
         b_prefix = prefix.encode('utf8')
         index = self.dct.follow_bytes(b_prefix, self.dct.ROOT)

dawg_python/wrapper.py

@@ -2,7 +2,6 @@
 from __future__ import absolute_import, unicode_literals
 import struct
 import array
-import pdb
 
 from . import units
 from .compat import int_from_byte
@@ -122,6 +121,7 @@ def start_edges(self, index, prefix=b""):
         index_stack. Otherwise, leave the stack empty, so next_edge() fails"""
 
         self.key = bytearray(prefix)
+        self.base_key_len = len(self.key)
         self._parent_index = index
         self._sib_index = None
         if self._guide.size():
@@ -133,6 +133,7 @@ def start_edges(self, index, prefix=b""):
                 if index is not None:
                     self._sib_index = next_index
                     self.key.append(child_label)
+                    self.decoded_key = self.key.decode('utf-8')
                     return True
 
     def next_edge(self):
@@ -146,9 +147,26 @@ def next_edge(self):
                                                 self._parent_index)
         if not self._sib_index:
             return False
-        
-        self.key.pop()
+
+        self.key = self.key[:self.base_key_len]
         self.key.append(sibling_label)
+        try:
+            self.decoded_key = self.key.decode('utf-8')
+        except UnicodeDecodeError:
+            #this sibling is multi-character. keep following its children til
+            #something is decodable
+            cur_index = self._sib_index
+            while True:
+                child_label = self._guide.child(self._sib_index)
+                cur_index = self._dic.follow_char(child_label, cur_index)
+                if not cur_index:
+                    return False
+                self.key.append(child_label)
+                try:
+                    self.decoded_key = self.key.decode('utf-8')
+                    break
+                except UnicodeDecodeError:
+                    pass
         return True
 
     def next(self):

tests/test_dawg.py

@@ -46,10 +46,20 @@ def test_keys(self):
         d = self.dawg()
         assert d.keys() == sorted(self.keys)
 
+    def test_edges(self):
+        d = self.dawg()
+        assert d.edges() == ['b', 'f']
+        assert d.edges('f') == ['fo']
+
     def test_iterkeys(self):
         d = self.dawg()
         assert list(d.iterkeys()) == d.keys()
 
+    def test_iter_edges(self):
+        d = self.dawg()
+        assert list(d.iteredges()) == ['b', 'f']
+        assert list(d.edges('f')) == ['fo']
+
     def test_completion(self):
         d = self.dawg()