Make smart futures avoid polling.
[python_utils.git] / dict_utils.py
1 #!/usr/bin/env python3
2
3 from itertools import islice
4 from typing import Any, Callable, Dict, Iterator, List, Tuple
5
6
7 def init_or_inc(
8     d: Dict[Any, Any],
9     key: Any,
10     *,
11     init_value: Any = 1,
12     inc_function: Callable[..., Any] = lambda x: x + 1
13 ) -> bool:
14     """
15     Initialize a dict value (if it doesn't exist) or increments it (using the
16     inc_function, which is customizable) if it already does exist.  Returns
17     True if the key already existed or False otherwise.
18
19     >>> d = {}
20     >>> init_or_inc(d, "test")
21     False
22     >>> init_or_inc(d, "test")
23     True
24     >>> init_or_inc(d, 'ing')
25     False
26     >>> d
27     {'test': 2, 'ing': 1}
28
29     """
30     if key in d.keys():
31         d[key] = inc_function(d[key])
32         return True
33     d[key] = init_value
34     return False
35
36
37 def shard(d: Dict[Any, Any], size: int) -> Iterator[Dict[Any, Any]]:
38     """
39     Shards a dict into N subdicts which, together, contain all keys/values
40     from the original unsharded dict.
41     """
42     items = d.items()
43     for x in range(0, len(d), size):
44         yield {key: value for (key, value) in islice(items, x, x + size)}
45
46
47 def coalesce_by_creating_list(key, new_value, old_value):
48     from list_utils import flatten
49     return flatten([new_value, old_value])
50
51
52 def coalesce_by_creating_set(key, new_value, old_value):
53     return set(coalesce_by_creating_list(key, new_value, old_value))
54
55
56 def coalesce_last_write_wins(key, new_value, old_value):
57     return new_value
58
59
60 def coalesce_first_write_wins(key, new_value, old_value):
61     return old_value
62
63
64 def raise_on_duplicated_keys(key, new_value, old_value):
65     raise Exception(f'Key {key} is duplicated in more than one input dict.')
66
67
68 def coalesce(
69         inputs: Iterator[Dict[Any, Any]],
70         *,
71         aggregation_function: Callable[[Any, Any], Any] = coalesce_by_creating_list
72 ) -> Dict[Any, Any]:
73     """Merge N dicts into one dict containing the union of all keys /
74     values in the input dicts.  When keys collide, apply the
75     aggregation_function which, by default, creates a list of values.
76     See also several other alternative functions for coalescing values
77     (coalesce_by_creating_set, coalesce_first_write_wins,
78     coalesce_last_write_wins, raise_on_duplicated_keys) or provide a
79     custom helper function.
80
81     >>> a = {'a': 1, 'b': 2}
82     >>> b = {'b': 1, 'c': 2, 'd': 3}
83     >>> c = {'c': 1, 'd': 2}
84     >>> coalesce([a, b, c])
85     {'a': 1, 'b': [1, 2], 'c': [1, 2], 'd': [2, 3]}
86
87     >>> coalesce([a, b, c], aggregation_function=coalesce_last_write_wins)
88     {'a': 1, 'b': 1, 'c': 1, 'd': 2}
89
90     >>> coalesce([a, b, c], aggregation_function=raise_on_duplicated_keys)
91     Traceback (most recent call last):
92     ...
93     Exception: Key b is duplicated in more than one input dict.
94
95     """
96     out: Dict[Any, Any] = {}
97     for d in inputs:
98         for key in d:
99             if key in out:
100                 value = aggregation_function(key, d[key], out[key])
101             else:
102                 value = d[key]
103             out[key] = value
104     return out
105
106
107 def item_with_max_value(d: Dict[Any, Any]) -> Tuple[Any, Any]:
108     """Returns the key and value with the max value in a dict.
109
110     >>> d = {'a': 1, 'b': 2, 'c': 3}
111     >>> item_with_max_value(d)
112     ('c', 3)
113     >>> item_with_max_value({})
114     Traceback (most recent call last):
115     ...
116     ValueError: max() arg is an empty sequence
117
118     """
119     return max(d.items(), key=lambda _: _[1])
120
121
122 def item_with_min_value(d: Dict[Any, Any]) -> Tuple[Any, Any]:
123     """Returns the key and value with the min value in a dict.
124
125     >>> d = {'a': 1, 'b': 2, 'c': 3}
126     >>> item_with_min_value(d)
127     ('a', 1)
128
129     """
130     return min(d.items(), key=lambda _: _[1])
131
132
133 def key_with_max_value(d: Dict[Any, Any]) -> Any:
134     """Returns the key with the max value in the dict.
135
136     >>> d = {'a': 1, 'b': 2, 'c': 3}
137     >>> key_with_max_value(d)
138     'c'
139
140     """
141     return item_with_max_value(d)[0]
142
143
144 def key_with_min_value(d: Dict[Any, Any]) -> Any:
145     """Returns the key with the min value in the dict.
146
147     >>> d = {'a': 1, 'b': 2, 'c': 3}
148     >>> key_with_min_value(d)
149     'a'
150
151     """
152     return item_with_min_value(d)[0]
153
154
155 def max_value(d: Dict[Any, Any]) -> Any:
156     """Returns the maximum value in the dict.
157
158     >>> d = {'a': 1, 'b': 2, 'c': 3}
159     >>> max_value(d)
160     3
161
162     """
163     return item_with_max_value(d)[1]
164
165
166 def min_value(d: Dict[Any, Any]) -> Any:
167     """Returns the minimum value in the dict.
168
169     >>> d = {'a': 1, 'b': 2, 'c': 3}
170     >>> min_value(d)
171     1
172
173     """
174     return item_with_min_value(d)[1]
175
176
177 def max_key(d: Dict[Any, Any]) -> Any:
178     """Returns the maximum key in dict (ignoring values totally)
179
180     >>> d = {'a': 3, 'b': 2, 'c': 1}
181     >>> max_key(d)
182     'c'
183
184     """
185     return max(d.keys())
186
187
188 def min_key(d: Dict[Any, Any]) -> Any:
189     """Returns the minimum key in dict (ignoring values totally)
190
191     >>> d = {'a': 3, 'b': 2, 'c': 1}
192     >>> min_key(d)
193     'a'
194
195     """
196     return min(d.keys())
197
198
199 def parallel_lists_to_dict(keys: List[Any], values: List[Any]) -> Dict[Any, Any]:
200     """Given two parallel lists (keys and values), create and return
201     a dict.
202
203     >>> k = ['name', 'phone', 'address', 'zip']
204     >>> v = ['scott', '555-1212', '123 main st.', '12345']
205     >>> parallel_lists_to_dict(k, v)
206     {'name': 'scott', 'phone': '555-1212', 'address': '123 main st.', 'zip': '12345'}
207
208     """
209     if len(keys) != len(values):
210         raise Exception("Parallel keys and values lists must have the same length")
211     return dict(zip(keys, values))
212
213
214 def dict_to_key_value_lists(d: Dict[Any, Any]) -> Tuple[List[Any], List[Any]]:
215     """
216     >>> d = {'name': 'scott', 'phone': '555-1212', 'address': '123 main st.', 'zip': '12345'}
217     >>> (k, v) = dict_to_key_value_lists(d)
218     >>> k
219     ['name', 'phone', 'address', 'zip']
220     >>> v
221     ['scott', '555-1212', '123 main st.', '12345']
222
223     """
224     r = ([], [])
225     for (k, v) in d.items():
226         r[0].append(k)
227         r[1].append(v)
228     return r
229
230
231 if __name__ == '__main__':
232     import doctest
233     doctest.testmod()