src/pyutils/collectionz/bidict.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """
   6 The :class:`pyutils.collectionz.bidict.BiDict` class is a subclass
   7 of :py:class:`dict` that implements a bidirectional dictionary.  That
   8 is, it maps each key to a value in constant time and each value back
   9 to the one or more keys it is associated with in constant time.  It
  10 does this by simply storing the data twice.
  11
  12 Sample usage::
  13
  14     # Initialize with a normal dict...
  15     third_party_wierdos = BiDict({
  16         'prometheus-fastapi-instrumentator': 'prometheus_fastapi_instrumentator',
  17         'scikit-learn': 'sklearn',
  18         'antlr4-python3-runtime' : 'antlr4',
  19         'python-dateutil': 'dateutil',
  20         'speechrecognition': 'speech_recognition',
  21         'beautifulsoup4': 'bs4',
  22         'python-dateutil': 'dateutil',
  23         'homeassistant-api': 'homeassistant_api',
  24     })
  25
  26     # Use in one direction:
  27     x = third_party_wierdos['scikit-learn']
  28
  29     # Use in opposite direction:
  30     y = third_party_wierdos.inverse['python_dateutil']
  31
  32     # Note: type(y) is List since one value may map back to multiple keys.
  33
  34 """
  35
  36
  37 class BiDict(dict):
  38     def __init__(self, *args, **kwargs):
  39         """
  40         A class that stores both a Mapping between keys and values and
  41         also the inverse mapping between values and their keys to
  42         allow for efficient lookups in either direction.  Because it
  43         is possible to have several keys with the same value, using
  44         the inverse map returns a sequence of keys.
  45
  46         >>> d = BiDict()
  47         >>> d['a'] = 1
  48         >>> d['b'] = 2
  49         >>> d['c'] = 2
  50         >>> d['a']
  51         1
  52         >>> d.inverse[1]
  53         ['a']
  54         >>> d.inverse[2]
  55         ['b', 'c']
  56         >>> len(d)
  57         3
  58         >>> del d['c']
  59         >>> len(d)
  60         2
  61         >>> d.inverse[2]
  62         ['b']
  63
  64         """
  65         super().__init__(*args, **kwargs)
  66         self.inverse = {}
  67         for key, value in self.items():
  68             self.inverse.setdefault(value, []).append(key)
  69
  70     def __setitem__(self, key, value):
  71         if key in self:
  72             old_value = self[key]
  73             self.inverse[old_value].remove(key)
  74         super().__setitem__(key, value)
  75         self.inverse.setdefault(value, []).append(key)
  76
  77     def __delitem__(self, key):
  78         value = self[key]
  79         self.inverse.setdefault(value, []).remove(key)
  80         if value in self.inverse and not self.inverse[value]:
  81             del self.inverse[value]
  82         super().__delitem__(key)
  83
  84
  85 if __name__ == '__main__':
  86     import doctest
  87
  88     doctest.testmod()