Move collections into collect space

[python_utils.git] / collect / bidict.py

diff --git a/collect/bidict.py b/collect/bidict.py
new file mode 100644 (file)
index 0000000..5ba3fc3
--- /dev/null
+++ b/collect/bidict.py
@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+
+class bidict(dict):
+    def __init__(self, *args, **kwargs):
+        super(bidict, self).__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)
+        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)