adding terminal indicators on edges; adding edge values for IntCompletionDawgs; adding tests for all

Author: EliFinkelshteyn

dawg_python/dawgs.py

@@ -149,13 +149,13 @@ def edges(self, prefix=""):
         if index is None:
             return res
 
-        completer = wrapper.Completer(self.dct, self.guide)
-        if not completer.start_edges(index, b_prefix):
+        edge_follower = wrapper.EdgeFollower(self.dct, self.guide)
+        if not edge_follower.start(index, b_prefix):
             return res
 
-        res.append(completer.decoded_key)
-        while completer.next_edge():
-            res.append(completer.decoded_key)
+        res.append(edge_follower.get_cur_edge())
+        while edge_follower.next():
+            res.append(edge_follower.get_cur_edge())
 
         return res
 
@@ -166,13 +166,13 @@ def iteredges(self, prefix=""):
         if index is None:
             return
 
-        completer = wrapper.Completer(self.dct, self.guide)
-        if not completer.start_edges(index, b_prefix):
+        edge_follower = wrapper.EdgeFollower(self.dct, self.guide)
+        if not edge_follower.start(index, b_prefix):
             return
 
-        yield completer.decoded_key
-        while completer.next_edge():
-            yield completer.decoded_key
+        yield edge_follower.get_cur_edge()
+        while edge_follower.next():
+            yield edge_follower.get_cur_edge()
 
     def iterkeys(self, prefix=""):
         b_prefix = prefix.encode('utf8')
@@ -312,15 +312,14 @@ def iterkeys(self, prefix=""):
             yield u_key
 
     def items(self, prefix=""):
+        index = self.dct.ROOT
         if not isinstance(prefix, bytes):
             prefix = prefix.encode('utf8')
-        res = []
-
-        index = self.dct.ROOT
         if prefix:
             index = self.dct.follow_bytes(prefix, index)
             if not index:
-                return res
+                return
+        res = []
 
         completer = wrapper.Completer(self.dct, self.guide)
         completer.start(index, prefix)
@@ -333,11 +332,30 @@ def items(self, prefix=""):
 
         return res
 
-    def iteritems(self, prefix=""):
+    def edges(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
+        res = []
+
+        edge_follower = wrapper.EdgeFollower(self.dct, self.guide)
+        if not edge_follower.start(index, prefix):
+            return res
+
+        res.append(edge_follower.decoded_key)
+        while edge_follower.next():
+            res.append(edge_follower.decoded_key)
 
+        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:
@@ -497,6 +515,43 @@ class IntCompletionDAWG(CompletionDAWG, IntDAWG):
     Dict-like class based on DAWG.
     It can store integer values for unicode keys and support key completion.
     """
+    def edges(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
+        res = []
+
+        edge_follower = wrapper.EdgeFollower(self.dct, self.guide)
+        if not edge_follower.start(index, prefix):
+            return res
+
+        res.append((edge_follower.decoded_key, edge_follower.value()))
+        while edge_follower.next():
+            res.append((edge_follower.decoded_key, edge_follower.value()))
+
+        return res
+
+    def iteredges(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
+
+        edge_follower = wrapper.EdgeFollower(self.dct, self.guide)
+        if not edge_follower.start(index, prefix):
+            return
+
+        yield (edge_follower.decoded_key, edge_follower.value())
+        while edge_follower.next():
+            yield (edge_follower.decoded_key, edge_follower.value())
+
     def items(self, prefix=""):
         if not isinstance(prefix, bytes):
             prefix = prefix.encode('utf8')

dawg_python/wrapper.py

@@ -17,29 +17,29 @@ def __init__(self):
     "Root index"
 
     def has_value(self, index):
-        "Checks if a given index is related to the end of a key."
+        #Checks if a given index is related to the end of a key.
         return units.has_leaf(self._units[index])
 
     def value(self, index):
-        "Gets a value from a given index."
+        #Gets a value from a given index.
         offset = units.offset(self._units[index])
         value_index = (index ^ offset) & units.PRECISION_MASK
         return units.value(self._units[value_index])
 
     def read(self, fp):
-        "Reads a dictionary from an input stream."
+        #Reads a dictionary from an input stream.
         base_size = struct.unpack(str("=I"), fp.read(4))[0]
         self._units.fromfile(fp, base_size)
 
     def contains(self, key):
-        "Exact matching."
+        #Exact matching.
         index = self.follow_bytes(key, self.ROOT)
         if index is None:
             return False
         return self.has_value(index)
 
     def find(self, key):
-        "Exact matching (returns value)"
+        #Exact matching (returns value)
         index = self.follow_bytes(key, self.ROOT)
         if index is None:
             return -1
@@ -48,7 +48,7 @@ def find(self, key):
         return self.value(index)
 
     def follow_char(self, label, index):
-        "Follows a transition"
+        #Follows a transition
         offset = units.offset(self._units[index])
         next_index = (index ^ offset ^ label) & units.PRECISION_MASK
 
@@ -58,7 +58,7 @@ def follow_char(self, label, index):
         return next_index
 
     def follow_bytes(self, s, index):
-        "Follows transitions."
+        #Follows transitions.
         for ch in s:
             index = self.follow_char(int_from_byte(ch), index)
             if index is None:
@@ -95,27 +95,17 @@ def size(self):
         return len(self._units)
 
 
-class Completer(object):
-
+class EdgeFollower(object):
     def __init__(self, dic=None, guide=None):
         self._dic = dic
         self._guide = guide
 
     def value(self):
-        return self._dic.value(self._last_index)
+        if self._dic.has_value(self._cur_index):
+            return self._dic.value(self._cur_index)
+        return False
 
     def start(self, index, prefix=b""):
-        "initial setup for a completer next() action on some prefix"
-
-        self.key = bytearray(prefix)
-
-        if self._guide.size():
-            self._index_stack = [index]
-            self._last_index = self._dic.ROOT
-        else:
-            self._index_stack = []
-
-    def start_edges(self, index, prefix=b""):
         """initial setup for a completer next_edge() action on some prefix. If
         there's a child for this prefix, we add that as the one item on the
         index_stack. Otherwise, leave the stack empty, so next_edge() fails"""
@@ -124,6 +114,7 @@ def start_edges(self, index, prefix=b""):
         self.base_key_len = len(self.key)
         self._parent_index = index
         self._sib_index = None
+        self._cur_index = None
         if self._guide.size():
             child_label = self._guide.child(index) # UCharType
 
@@ -132,19 +123,21 @@ def start_edges(self, index, prefix=b""):
                 next_index = self._dic.follow_char(child_label, index)
                 if index is not None:
                     self._sib_index = next_index
+                    self._cur_index = self._sib_index
                     self.key.append(child_label)
                     self.decoded_key = self.key.decode('utf-8')
                     return True
 
-    def next_edge(self):
-        "Gets the next edge (not necessarily a terminal)"
+    def next(self):
+        #Gets the next edge (not necessarily a terminal)
 
         if not self._sib_index:
             return False
 
         sibling_label = self._guide.sibling(self._sib_index)
         self._sib_index = self._dic.follow_char(sibling_label,
                                                 self._parent_index)
+        self._cur_index = self._sib_index
         if not self._sib_index:
             return False
 
@@ -153,13 +146,13 @@ def next_edge(self):
         try:
             self.decoded_key = self.key.decode('utf-8')
         except UnicodeDecodeError:
-            #this sibling is multi-character. keep following its children til
+            #this sibling is a multibyte char. 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:
+                self._cur_index = self._dic.follow_char(child_label,
+                                                        self._cur_index)
+                if not self._cur_index:
                     return False
                 self.key.append(child_label)
                 try:
@@ -169,8 +162,32 @@ def next_edge(self):
                     pass
         return True
 
+    def get_cur_edge(self):
+        return (self.decoded_key, self._dic.has_value(self._cur_index))
+
+
+class Completer(object):
+
+    def __init__(self, dic=None, guide=None):
+        self._dic = dic
+        self._guide = guide
+
+    def value(self):
+        return self._dic.value(self._last_index)
+
+    def start(self, index, prefix=b""):
+        #initial setup for a completer next() action on some prefix
+
+        self.key = bytearray(prefix)
+
+        if self._guide.size():
+            self._index_stack = [index]
+            self._last_index = self._dic.ROOT
+        else:
+            self._index_stack = []
+
     def next(self):
-        "Gets the next key"
+        #Gets the next key
 
         if not self._index_stack:
             return False

tests/test_dawg.py

@@ -48,17 +48,19 @@ def test_keys(self):
 
     def test_edges(self):
         d = self.dawg()
-        assert d.edges() == ['b', 'f']
-        assert d.edges('f') == ['fo']
+        assert d.edges() == [('b', False), ('f', True)]
+        assert d.edges('b') == [('ba', False)]
+        assert d.edges('fo') == [('foo', True)]
 
     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']
+        assert list(d.iteredges()) == [('b', False), ('f', True)]
+        assert list(d.iteredges('b')) == [('ba', False)]
+        assert list(d.edges('fo')) == [('foo', True)]
 
     def test_completion(self):
         d = self.dawg()
@@ -129,3 +131,11 @@ def test_completion_keys_with_prefix(self):
 
     def test_completion_items(self):
         assert self.dawg().items() == sorted(self.payload.items(), key=lambda r: r[0])
+
+    def test_completion_edges(self):
+        assert self.dawg().edges('ba') == [('bar', 5)]
+        assert self.dawg().edges('foob') == [('fooba', False)]
+
+    def test_completion_iteredges(self):
+        assert list(self.dawg().iteredges('ba')) == [('bar', 5)]
+        assert list(self.dawg().iteredges('foob')) == [('fooba', False)]