X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=collect%2Fbidict.py;h=8153e54446c8a4c5a36f96ba85e0b93c9afafe18;hb=ed8fa2b10b0177b15b7423263bdd390efde2f0c8;hp=5ba3fc30f6780537509ba26f49202852ef8a5e05;hpb=0bc6e4312cad0f997751739e750954ac39dfa6cc;p=python_utils.git

diff --git a/collect/bidict.py b/collect/bidict.py
index 5ba3fc3..8153e54 100644
--- a/collect/bidict.py
+++ b/collect/bidict.py
@@ -1,20 +1,53 @@
 #!/usr/bin/env python3
 
-class bidict(dict):
+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)
 
     def __setitem__(self, key, value):
         if key in self:
-            self.inverse[self[key]].remove(key)
-        super(bidict, self).__setitem__(key, value)
+            old_value = self[key]
+            self.inverse[old_value].remove(key)
+        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)
+        value = self[key]
+        self.inverse.setdefault(value, []).remove(key)
+        if value in self.inverse and not self.inverse[value]:
+            del self.inverse[value]
+        super().__delitem__(key)
+
+
+if __name__ == '__main__':
+    import doctest
+    doctest.testmod()