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