Add doctests to some of this stuff.
[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, 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     if key in d.keys():
30         d[key] = inc_function(d[key])
31         return True
32     d[key] = init_value
33     return False
34
35
36 def shard(d: Dict[Any, Any], size: int) -> Iterator[Dict[Any, Any]]:
37     """
38     Shards a dict into N subdicts which, together, contain all keys/values
39     from the original unsharded dict.
40     """
41     items = d.items()
42     for x in range(0, len(d), size):
43         yield {key: value for (key, value) in islice(items, x, x + size)}
44
45
46 def coalesce_by_creating_list(key, v1, v2):
47     from list_utils import flatten
48     return flatten([v1, v2])
49
50
51 def coalesce_by_creating_set(key, v1, v2):
52     return set(coalesce_by_creating_list(key, v1, v2))
53
54
55 def raise_on_duplicated_keys(key, v1, v2):
56     raise Exception(f'Key {key} is duplicated in more than one input dict.')
57
58
59 def coalesce(
60         inputs: Iterator[Dict[Any, Any]],
61         *,
62         aggregation_function: Callable[[Any, Any], Any] = coalesce_by_creating_list
63 ) -> Dict[Any, Any]:
64     """Merge N dicts into one dict containing the union of all keys/values in
65     the input dicts.  When keys collide, apply the aggregation_function which,
66     by default, creates a list of values.  See also coalesce_by_creating_set or
67     provide a user defined aggregation_function.
68
69     >>> a = {'a': 1, 'b': 2}
70     >>> b = {'b': 1, 'c': 2, 'd': 3}
71     >>> c = {'c': 1, 'd': 2}
72     >>> coalesce([a, b, c])
73     {'a': 1, 'b': [1, 2], 'c': [1, 2], 'd': [2, 3]}
74     """
75     out: Dict[Any, Any] = {}
76     for d in inputs:
77         for key in d:
78             if key in out:
79                 value = aggregation_function(key, d[key], out[key])
80             else:
81                 value = d[key]
82             out[key] = value
83     return out
84
85
86 def item_with_max_value(d: Dict[Any, Any]) -> Tuple[Any, Any]:
87     """Returns the key and value with the max value in a dict.
88
89     >>> d = {'a': 1, 'b': 2, 'c': 3}
90     >>> item_with_max_value(d)
91     ('c', 3)
92     >>> item_with_max_value({})
93     Traceback (most recent call last):
94     ...
95     ValueError: max() arg is an empty sequence
96     """
97     return max(d.items(), key=lambda _: _[1])
98
99
100 def item_with_min_value(d: Dict[Any, Any]) -> Tuple[Any, Any]:
101     """Returns the key and value with the min value in a dict.
102
103     >>> d = {'a': 1, 'b': 2, 'c': 3}
104     >>> item_with_min_value(d)
105     ('a', 1)
106     """
107     return min(d.items(), key=lambda _: _[1])
108
109
110 def key_with_max_value(d: Dict[Any, Any]) -> Any:
111     """Returns the key with the max value in the dict.
112
113     >>> d = {'a': 1, 'b': 2, 'c': 3}
114     >>> key_with_max_value(d)
115     'c'
116     """
117     return item_with_max_value(d)[0]
118
119
120 def key_with_min_value(d: Dict[Any, Any]) -> Any:
121     """Returns the key with the min value in the dict.
122
123     >>> d = {'a': 1, 'b': 2, 'c': 3}
124     >>> key_with_min_value(d)
125     'a'
126     """
127     return item_with_min_value(d)[0]
128
129
130 def max_value(d: Dict[Any, Any]) -> Any:
131     """Returns the maximum value in the dict.
132
133     >>> d = {'a': 1, 'b': 2, 'c': 3}
134     >>> max_value(d)
135     3
136     """
137     return item_with_max_value(d)[1]
138
139
140 def min_value(d: Dict[Any, Any]) -> Any:
141     """Returns the minimum value in the dict.
142
143     >>> d = {'a': 1, 'b': 2, 'c': 3}
144     >>> min_value(d)
145     1
146     """
147     return item_with_min_value(d)[1]
148
149
150 def max_key(d: Dict[Any, Any]) -> Any:
151     """Returns the maximum key in dict (ignoring values totally)
152
153     >>> d = {'a': 3, 'b': 2, 'c': 1}
154     >>> max_key(d)
155     'c'
156     """
157     return max(d.keys())
158
159
160 def min_key(d: Dict[Any, Any]) -> Any:
161     """Returns the minimum key in dict (ignoring values totally)
162
163     >>> d = {'a': 3, 'b': 2, 'c': 1}
164     >>> min_key(d)
165     'a'
166     """
167     return min(d.keys())
168
169
170 if __name__ == '__main__':
171     import doctest
172     doctest.testmod()