Add tests on bidict. Fix bug in string_utils ngrams.

1 files changed, 34 insertions, 3 deletions

diff --git a/collect/bidict.py b/collect/bidict.py
index 5ba3fc3..e162179 100644
--- a/collect/bidict.py
+++ b/collect/bidict.py
@@ -2,7 +2,33 @@
 
 class bidict(dict):
     def __init__(self, *args, **kwargs):
-        super(bidict, self).__init__(*args, **kwargs)
+        """
+        A class that stores both a Mapping between keys and values and
+        also the inverse mapping between values and their keys to
+        allow for efficient lookups in either direction.  Because it
+        is possible to have several keys with the same value, using
+        the inverse map returns a sequence of keys.
+
+        >>> d = bidict()
+        >>> d['a'] = 1
+        >>> d['b'] = 2
+        >>> d['c'] = 2
+        >>> d['a']
+        1
+        >>> d.inverse[1]
+        ['a']
+        >>> d.inverse[2]
+        ['b', 'c']
+        >>> len(d)
+        3
+        >>> del d['c']
+        >>> len(d)
+        2
+        >>> d.inverse[2]
+        ['b']
+
+        """
+        super().__init__(*args, **kwargs)
         self.inverse = {}
         for key, value in self.items():
             self.inverse.setdefault(value, []).append(key)
@@ -10,11 +36,16 @@ class bidict(dict):
     def __setitem__(self, key, value):
         if key in self:
             self.inverse[self[key]].remove(key)
-        super(bidict, self).__setitem__(key, value)
+        super().__setitem__(key, value)
         self.inverse.setdefault(value, []).append(key)
 
     def __delitem__(self, key):
         self.inverse.setdefault(self[key], []).remove(key)
         if self[key] in self.inverse and not self.inverse[self[key]]:
             del self.inverse[self[key]]
-        super(bidict, self).__delitem__(key)
+        super().__delitem__(key)
+
+
+if __name__ == '__main__':
+    import doctest
+    doctest.testmod()