summaryrefslogtreecommitdiff
path: root/dict_utils.py
blob: 451a87dadf08d8632ac6f593dfb592116a05779b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
#!/usr/bin/env python3

from itertools import islice
from typing import Any, Callable, Dict, Iterator, List, Tuple


def init_or_inc(
    d: Dict[Any, Any],
    key: Any,
    *,
    init_value: Any = 1,
    inc_function: Callable[..., Any] = lambda x: x + 1,
) -> bool:
    """
    Initialize a dict value (if it doesn't exist) or increments it (using the
    inc_function, which is customizable) if it already does exist.  Returns
    True if the key already existed or False otherwise.

    >>> d = {}
    >>> init_or_inc(d, "test")
    False
    >>> init_or_inc(d, "test")
    True
    >>> init_or_inc(d, 'ing')
    False
    >>> d
    {'test': 2, 'ing': 1}

    """
    if key in d.keys():
        d[key] = inc_function(d[key])
        return True
    d[key] = init_value
    return False


def shard(d: Dict[Any, Any], size: int) -> Iterator[Dict[Any, Any]]:
    """
    Shards a dict into N subdicts which, together, contain all keys/values
    from the original unsharded dict.

    """
    items = d.items()
    for x in range(0, len(d), size):
        yield {key: value for (key, value) in islice(items, x, x + size)}


def coalesce_by_creating_list(key, new_value, old_value):
    from list_utils import flatten

    return flatten([new_value, old_value])


def coalesce_by_creating_set(key, new_value, old_value):
    return set(coalesce_by_creating_list(key, new_value, old_value))


def coalesce_last_write_wins(key, new_value, old_value):
    return new_value


def coalesce_first_write_wins(key, new_value, old_value):
    return old_value


def raise_on_duplicated_keys(key, new_value, old_value):
    raise Exception(f'Key {key} is duplicated in more than one input dict.')


def coalesce(
    inputs: Iterator[Dict[Any, Any]],
    *,
    aggregation_function: Callable[[Any, Any, Any], Any] = coalesce_by_creating_list,
) -> Dict[Any, Any]:
    """Merge N dicts into one dict containing the union of all keys /
    values in the input dicts.  When keys collide, apply the
    aggregation_function which, by default, creates a list of values.
    See also several other alternative functions for coalescing values
    (coalesce_by_creating_set, coalesce_first_write_wins,
    coalesce_last_write_wins, raise_on_duplicated_keys) or provide a
    custom helper function.

    >>> a = {'a': 1, 'b': 2}
    >>> b = {'b': 1, 'c': 2, 'd': 3}
    >>> c = {'c': 1, 'd': 2}
    >>> coalesce([a, b, c])
    {'a': 1, 'b': [1, 2], 'c': [1, 2], 'd': [2, 3]}

    >>> coalesce([a, b, c], aggregation_function=coalesce_last_write_wins)
    {'a': 1, 'b': 1, 'c': 1, 'd': 2}

    >>> coalesce([a, b, c], aggregation_function=raise_on_duplicated_keys)
    Traceback (most recent call last):
    ...
    Exception: Key b is duplicated in more than one input dict.

    """
    out: Dict[Any, Any] = {}
    for d in inputs:
        for key in d:
            if key in out:
                value = aggregation_function(key, d[key], out[key])
            else:
                value = d[key]
            out[key] = value
    return out


def item_with_max_value(d: Dict[Any, Any]) -> Tuple[Any, Any]:
    """Returns the key and value with the max value in a dict.

    >>> d = {'a': 1, 'b': 2, 'c': 3}
    >>> item_with_max_value(d)
    ('c', 3)
    >>> item_with_max_value({})
    Traceback (most recent call last):
    ...
    ValueError: max() arg is an empty sequence

    """
    return max(d.items(), key=lambda _: _[1])


def item_with_min_value(d: Dict[Any, Any]) -> Tuple[Any, Any]:
    """Returns the key and value with the min value in a dict.

    >>> d = {'a': 1, 'b': 2, 'c': 3}
    >>> item_with_min_value(d)
    ('a', 1)

    """
    return min(d.items(), key=lambda _: _[1])


def key_with_max_value(d: Dict[Any, Any]) -> Any:
    """Returns the key with the max value in the dict.

    >>> d = {'a': 1, 'b': 2, 'c': 3}
    >>> key_with_max_value(d)
    'c'

    """
    return item_with_max_value(d)[0]


def key_with_min_value(d: Dict[Any, Any]) -> Any:
    """Returns the key with the min value in the dict.

    >>> d = {'a': 1, 'b': 2, 'c': 3}
    >>> key_with_min_value(d)
    'a'

    """
    return item_with_min_value(d)[0]


def max_value(d: Dict[Any, Any]) -> Any:
    """Returns the maximum value in the dict.

    >>> d = {'a': 1, 'b': 2, 'c': 3}
    >>> max_value(d)
    3

    """
    return item_with_max_value(d)[1]


def min_value(d: Dict[Any, Any]) -> Any:
    """Returns the minimum value in the dict.

    >>> d = {'a': 1, 'b': 2, 'c': 3}
    >>> min_value(d)
    1

    """
    return item_with_min_value(d)[1]


def max_key(d: Dict[Any, Any]) -> Any:
    """Returns the maximum key in dict (ignoring values totally)

    >>> d = {'a': 3, 'b': 2, 'c': 1}
    >>> max_key(d)
    'c'

    """
    return max(d.keys())


def min_key(d: Dict[Any, Any]) -> Any:
    """Returns the minimum key in dict (ignoring values totally)

    >>> d = {'a': 3, 'b': 2, 'c': 1}
    >>> min_key(d)
    'a'

    """
    return min(d.keys())


def parallel_lists_to_dict(keys: List[Any], values: List[Any]) -> Dict[Any, Any]:
    """Given two parallel lists (keys and values), create and return
    a dict.

    >>> k = ['name', 'phone', 'address', 'zip']
    >>> v = ['scott', '555-1212', '123 main st.', '12345']
    >>> parallel_lists_to_dict(k, v)
    {'name': 'scott', 'phone': '555-1212', 'address': '123 main st.', 'zip': '12345'}

    """
    if len(keys) != len(values):
        raise Exception("Parallel keys and values lists must have the same length")
    return dict(zip(keys, values))


def dict_to_key_value_lists(d: Dict[Any, Any]) -> Tuple[List[Any], List[Any]]:
    """
    >>> d = {'name': 'scott', 'phone': '555-1212', 'address': '123 main st.', 'zip': '12345'}
    >>> (k, v) = dict_to_key_value_lists(d)
    >>> k
    ['name', 'phone', 'address', 'zip']
    >>> v
    ['scott', '555-1212', '123 main st.', '12345']

    """
    r: Tuple[List[Any], List[Any]] = ([], [])
    for (k, v) in d.items():
        r[0].append(k)
        r[1].append(v)
    return r


if __name__ == '__main__':
    import doctest

    doctest.testmod()