Since this thing is on the innerwebs I suppose it should have a

[python_utils.git] / collect / bidict.py

diff --git a/collect/bidict.py b/collect/bidict.py
index 5ba3fc30f6780537509ba26f49202852ef8a5e05..375721e4a2652564ddfe42c6dc61925f8c7b4a76 100644 (file)
--- a/collect/bidict.py
+++ b/collect/bidict.py
@@ -1,20 +1,59 @@
 #!/usr/bin/env python3
 
-class bidict(dict):
+# © Copyright 2021-2022, Scott Gasch
+
+"""Bidirectional Dictionary."""
+
+
+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()