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