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             self.inverse[self[key]].remove(key)
  39         super().__setitem__(key, value)
  40         self.inverse.setdefault(value, []).append(key)
  41
  42     def __delitem__(self, key):
  43         self.inverse.setdefault(self[key], []).remove(key)
  44         if self[key] in self.inverse and not self.inverse[self[key]]:
  45             del self.inverse[self[key]]
  46         super().__delitem__(key)
  47
  48
  49 if __name__ == '__main__':
  50     import doctest
  51     doctest.testmod()