collect/bidict.py

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