collect/bidict.py

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