collect/bidict.py

   1 #!/usr/bin/env python3
   2
   3 class BiDict(dict):
   4     def __init__(self, *args, **kwargs):
   5         """
   6         A class that stores both a Mapping between keys and values and
   7         also the inverse mapping between values and their keys to
   8         allow for efficient lookups in either direction.  Because it
   9         is possible to have several keys with the same value, using
  10         the inverse map returns a sequence of keys.
  11
  12         >>> d = BiDict()
  13         >>> d['a'] = 1
  14         >>> d['b'] = 2
  15         >>> d['c'] = 2
  16         >>> d['a']
  17         1
  18         >>> d.inverse[1]
  19         ['a']
  20         >>> d.inverse[2]
  21         ['b', 'c']
  22         >>> len(d)
  23         3
  24         >>> del d['c']
  25         >>> len(d)
  26         2
  27         >>> d.inverse[2]
  28         ['b']
  29
  30         """
  31         super().__init__(*args, **kwargs)
  32         self.inverse = {}
  33         for key, value in self.items():
  34             self.inverse.setdefault(value, []).append(key)
  35
  36     def __setitem__(self, key, value):
  37         if key in self:
  38             old_value = self[key]
  39             self.inverse[old_value].remove(key)
  40         super().__setitem__(key, value)
  41         self.inverse.setdefault(value, []).append(key)
  42
  43     def __delitem__(self, key):
  44         value = self[key]
  45         self.inverse.setdefault(value, []).remove(key)
  46         if value in self.inverse and not self.inverse[value]:
  47             del self.inverse[value]
  48         super().__delitem__(key)
  49
  50
  51 if __name__ == '__main__':
  52     import doctest
  53     doctest.testmod()