| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869 |
- # encoding: utf-8
- """Miscellaneous context managers."""
- from __future__ import annotations
- import warnings
- from types import TracebackType
- from typing import Any
- # Copyright (c) IPython Development Team.
- # Distributed under the terms of the Modified BSD License.
- class preserve_keys:
- """Preserve a set of keys in a dictionary.
- Upon entering the context manager the current values of the keys
- will be saved. Upon exiting, the dictionary will be updated to
- restore the original value of the preserved keys. Preserved keys
- which did not exist when entering the context manager will be
- deleted.
- Examples
- --------
- >>> d = {'a': 1, 'b': 2, 'c': 3}
- >>> with preserve_keys(d, 'b', 'c', 'd'):
- ... del d['a']
- ... del d['b'] # will be reset to 2
- ... d['c'] = None # will be reset to 3
- ... d['d'] = 4 # will be deleted
- ... d['e'] = 5
- ... print(sorted(d.items()))
- ...
- [('c', None), ('d', 4), ('e', 5)]
- >>> print(sorted(d.items()))
- [('b', 2), ('c', 3), ('e', 5)]
- """
- def __init__(self, dictionary: dict[Any, Any], *keys: Any) -> None:
- self.dictionary = dictionary
- self.keys = keys
- def __enter__(self) -> None:
- # Actions to perform upon exiting.
- to_delete: list[Any] = []
- to_update: dict[Any, Any] = {}
- d = self.dictionary
- for k in self.keys:
- if k in d:
- to_update[k] = d[k]
- else:
- to_delete.append(k)
- self.to_delete = to_delete
- self.to_update = to_update
- def __exit__(
- self,
- exc_type: type[BaseException] | None,
- exc_val: BaseException | None,
- exc_tb: TracebackType | None,
- ) -> None:
- d = self.dictionary
- for k in self.to_delete:
- d.pop(k, None)
- d.update(self.to_update)
|
