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 e16217994334a6af750cea13bd88a03f309b2329..375721e4a2652564ddfe42c6dc61925f8c7b4a76 100644 (file)
--- a/collect/bidict.py
+++ b/collect/bidict.py
@@ -1,6 +1,11 @@
 #!/usr/bin/env python3
 
-class bidict(dict):
+# © Copyright 2021-2022, Scott Gasch
+
+"""Bidirectional Dictionary."""
+
+
+class BiDict(dict):
     def __init__(self, *args, **kwargs):
         """
         A class that stores both a Mapping between keys and values and
@@ -9,7 +14,7 @@ class bidict(dict):
         is possible to have several keys with the same value, using
         the inverse map returns a sequence of keys.
 
-        >>> d = bidict()
+        >>> d = BiDict()
         >>> d['a'] = 1
         >>> d['b'] = 2
         >>> d['c'] = 2
@@ -35,17 +40,20 @@ class bidict(dict):
 
     def __setitem__(self, key, value):
         if key in self:
-            self.inverse[self[key]].remove(key)
+            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]]
+        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()