1  # encoding: utf-8
   2  """Miscellaneous context managers.
   3  """
   4  
   5  import warnings
   6  
   7  # Copyright (c) IPython Development Team.
   8  # Distributed under the terms of the Modified BSD License.
   9  
  10  
  11  class preserve_keys(object):
  12      """Preserve a set of keys in a dictionary.
  13  
  14      Upon entering the context manager the current values of the keys
  15      will be saved. Upon exiting, the dictionary will be updated to
  16      restore the original value of the preserved keys. Preserved keys
  17      which did not exist when entering the context manager will be
  18      deleted.
  19  
  20      Examples
  21      --------
  22  
  23      >>> d = {'a': 1, 'b': 2, 'c': 3}
  24      >>> with preserve_keys(d, 'b', 'c', 'd'):
  25      ...     del d['a']
  26      ...     del d['b']      # will be reset to 2
  27      ...     d['c'] = None   # will be reset to 3
  28      ...     d['d'] = 4      # will be deleted
  29      ...     d['e'] = 5
  30      ...     print(sorted(d.items()))
  31      ...
  32      [('c', None), ('d', 4), ('e', 5)]
  33      >>> print(sorted(d.items()))
  34      [('b', 2), ('c', 3), ('e', 5)]
  35      """
  36  
  37      def __init__(self, dictionary, *keys):
  38          self.dictionary = dictionary
  39          self.keys = keys
  40  
  41      def __enter__(self):
  42          # Actions to perform upon exiting.
  43          to_delete = []
  44          to_update = {}
  45  
  46          d = self.dictionary
  47          for k in self.keys:
  48              if k in d:
  49                  to_update[k] = d[k]
  50              else:
  51                  to_delete.append(k)
  52  
  53          self.to_delete = to_delete
  54          self.to_update = to_update
  55  
  56      def __exit__(self, *exc_info):
  57          d = self.dictionary
  58  
  59          for k in self.to_delete:
  60              d.pop(k, None)
  61          d.update(self.to_update)
  62