|
1 | 1 | import collections |
| 2 | +import functools |
2 | 3 | import itertools |
| 4 | +import operator |
3 | 5 |
|
4 | 6 |
|
5 | 7 | # from jaraco.collections 3.5.1 |
@@ -54,3 +56,139 @@ def __contains__(self, other): |
54 | 56 |
|
55 | 57 | def __len__(self): |
56 | 58 | return len(list(iter(self))) |
| 59 | + |
| 60 | + |
| 61 | +# from jaraco.collections 3.7 |
| 62 | +class RangeMap(dict): |
| 63 | + """ |
| 64 | + A dictionary-like object that uses the keys as bounds for a range. |
| 65 | + Inclusion of the value for that range is determined by the |
| 66 | + key_match_comparator, which defaults to less-than-or-equal. |
| 67 | + A value is returned for a key if it is the first key that matches in |
| 68 | + the sorted list of keys. |
| 69 | +
|
| 70 | + One may supply keyword parameters to be passed to the sort function used |
| 71 | + to sort keys (i.e. key, reverse) as sort_params. |
| 72 | +
|
| 73 | + Let's create a map that maps 1-3 -> 'a', 4-6 -> 'b' |
| 74 | +
|
| 75 | + >>> r = RangeMap({3: 'a', 6: 'b'}) # boy, that was easy |
| 76 | + >>> r[1], r[2], r[3], r[4], r[5], r[6] |
| 77 | + ('a', 'a', 'a', 'b', 'b', 'b') |
| 78 | +
|
| 79 | + Even float values should work so long as the comparison operator |
| 80 | + supports it. |
| 81 | +
|
| 82 | + >>> r[4.5] |
| 83 | + 'b' |
| 84 | +
|
| 85 | + But you'll notice that the way rangemap is defined, it must be open-ended |
| 86 | + on one side. |
| 87 | +
|
| 88 | + >>> r[0] |
| 89 | + 'a' |
| 90 | + >>> r[-1] |
| 91 | + 'a' |
| 92 | +
|
| 93 | + One can close the open-end of the RangeMap by using undefined_value |
| 94 | +
|
| 95 | + >>> r = RangeMap({0: RangeMap.undefined_value, 3: 'a', 6: 'b'}) |
| 96 | + >>> r[0] |
| 97 | + Traceback (most recent call last): |
| 98 | + ... |
| 99 | + KeyError: 0 |
| 100 | +
|
| 101 | + One can get the first or last elements in the range by using RangeMap.Item |
| 102 | +
|
| 103 | + >>> last_item = RangeMap.Item(-1) |
| 104 | + >>> r[last_item] |
| 105 | + 'b' |
| 106 | +
|
| 107 | + .last_item is a shortcut for Item(-1) |
| 108 | +
|
| 109 | + >>> r[RangeMap.last_item] |
| 110 | + 'b' |
| 111 | +
|
| 112 | + Sometimes it's useful to find the bounds for a RangeMap |
| 113 | +
|
| 114 | + >>> r.bounds() |
| 115 | + (0, 6) |
| 116 | +
|
| 117 | + RangeMap supports .get(key, default) |
| 118 | +
|
| 119 | + >>> r.get(0, 'not found') |
| 120 | + 'not found' |
| 121 | +
|
| 122 | + >>> r.get(7, 'not found') |
| 123 | + 'not found' |
| 124 | +
|
| 125 | + One often wishes to define the ranges by their left-most values, |
| 126 | + which requires use of sort params and a key_match_comparator. |
| 127 | +
|
| 128 | + >>> r = RangeMap({1: 'a', 4: 'b'}, |
| 129 | + ... sort_params=dict(reverse=True), |
| 130 | + ... key_match_comparator=operator.ge) |
| 131 | + >>> r[1], r[2], r[3], r[4], r[5], r[6] |
| 132 | + ('a', 'a', 'a', 'b', 'b', 'b') |
| 133 | +
|
| 134 | + That wasn't nearly as easy as before, so an alternate constructor |
| 135 | + is provided: |
| 136 | +
|
| 137 | + >>> r = RangeMap.left({1: 'a', 4: 'b', 7: RangeMap.undefined_value}) |
| 138 | + >>> r[1], r[2], r[3], r[4], r[5], r[6] |
| 139 | + ('a', 'a', 'a', 'b', 'b', 'b') |
| 140 | +
|
| 141 | + """ |
| 142 | + |
| 143 | + def __init__(self, source, sort_params={}, key_match_comparator=operator.le): |
| 144 | + dict.__init__(self, source) |
| 145 | + self.sort_params = sort_params |
| 146 | + self.match = key_match_comparator |
| 147 | + |
| 148 | + @classmethod |
| 149 | + def left(cls, source): |
| 150 | + return cls( |
| 151 | + source, sort_params=dict(reverse=True), key_match_comparator=operator.ge |
| 152 | + ) |
| 153 | + |
| 154 | + def __getitem__(self, item): |
| 155 | + sorted_keys = sorted(self.keys(), **self.sort_params) |
| 156 | + if isinstance(item, RangeMap.Item): |
| 157 | + result = self.__getitem__(sorted_keys[item]) |
| 158 | + else: |
| 159 | + key = self._find_first_match_(sorted_keys, item) |
| 160 | + result = dict.__getitem__(self, key) |
| 161 | + if result is RangeMap.undefined_value: |
| 162 | + raise KeyError(key) |
| 163 | + return result |
| 164 | + |
| 165 | + def get(self, key, default=None): |
| 166 | + """ |
| 167 | + Return the value for key if key is in the dictionary, else default. |
| 168 | + If default is not given, it defaults to None, so that this method |
| 169 | + never raises a KeyError. |
| 170 | + """ |
| 171 | + try: |
| 172 | + return self[key] |
| 173 | + except KeyError: |
| 174 | + return default |
| 175 | + |
| 176 | + def _find_first_match_(self, keys, item): |
| 177 | + is_match = functools.partial(self.match, item) |
| 178 | + matches = list(filter(is_match, keys)) |
| 179 | + if matches: |
| 180 | + return matches[0] |
| 181 | + raise KeyError(item) |
| 182 | + |
| 183 | + def bounds(self): |
| 184 | + sorted_keys = sorted(self.keys(), **self.sort_params) |
| 185 | + return (sorted_keys[RangeMap.first_item], sorted_keys[RangeMap.last_item]) |
| 186 | + |
| 187 | + # some special values for the RangeMap |
| 188 | + undefined_value = type(str('RangeValueUndefined'), (), {})() |
| 189 | + |
| 190 | + class Item(int): |
| 191 | + "RangeMap Item" |
| 192 | + |
| 193 | + first_item = Item(0) |
| 194 | + last_item = Item(-1) |
0 commit comments