sorteddict.py revision 12563:8d59ed22ae79 
1# Copyright (c) 2006-2009 Nathan Binkert <nate@binkert.org>
2# All rights reserved.
3#
4# Redistribution and use in source and binary forms, with or without
5# modification, are permitted provided that the following conditions are
6# met: redistributions of source code must retain the above copyright
7# notice, this list of conditions and the following disclaimer;
8# redistributions in binary form must reproduce the above copyright
9# notice, this list of conditions and the following disclaimer in the
10# documentation and/or other materials provided with the distribution;
11# neither the name of the copyright holders nor the names of its
12# contributors may be used to endorse or promote products derived from
13# this software without specific prior written permission.
14#
15# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
16# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
17# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
18# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
19# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
20# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
21# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
22# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
23# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26
27from __future__ import print_function
28
29from bisect import bisect_left, bisect_right
30
31class SortedDict(dict):
32    def _get_sorted(self):
33        return getattr(self, '_sorted', sorted)
34    def _set_sorted(self, val):
35        self._sorted = val
36        self._del_keys()
37    sorted = property(_get_sorted, _set_sorted)
38
39    @property
40    def _keys(self):
41        try:
42            return self._sorted_keys
43        except AttributeError:
44            _sorted_keys = self.sorted(dict.iterkeys(self))
45            self._sorted_keys = _sorted_keys
46            return _sorted_keys
47
48    def _left_eq(self, key):
49        index = self._left_ge(self, key)
50        if self._keys[index] != key:
51            raise KeyError(key)
52        return index
53
54    def _right_eq(self, key):
55        index = self._right_le(self, key)
56        if self._keys[index] != key:
57            raise KeyError(key)
58        return index
59
60    def _right_lt(self, key):
61        index = bisect_left(self._keys, key)
62        if index:
63            return index - 1
64        raise KeyError(key)
65
66    def _right_le(self, key):
67        index = bisect_right(self._keys, key)
68        if index:
69            return index - 1
70        raise KeyError(key)
71
72    def _left_gt(self, key):
73        index = bisect_right(self._keys, key)
74        if index != len(self._keys):
75            return index
76        raise KeyError(key)
77
78    def _left_ge(self, key):
79        index = bisect_left(self._keys, key)
80        if index != len(self._keys):
81            return index
82        raise KeyError(key)
83
84    def _del_keys(self):
85        try:
86            del self._sorted_keys
87        except AttributeError:
88            pass
89
90    def __repr__(self):
91        return 'SortedDict({%s})' % ', '.join('%r: %r' % item
92                                              for item in self.iteritems())
93    def __setitem__(self, key, item):
94        dict.__setitem__(self, key, item)
95        self._del_keys()
96
97    def __delitem__(self, key):
98        dict.__delitem__(self, key)
99        self._del_keys()
100
101    def clear(self):
102        self.data.clear()
103        self._del_keys()
104
105    def copy(self):
106        t = type(self)
107        return t(self)
108
109    def keys(self):
110        return self._keys[:]
111
112    def values(self):
113        return list(self.itervalues())
114
115    def items(self):
116        return list(self.iteritems())
117
118    def iterkeys(self):
119        return iter(self._keys)
120
121    def itervalues(self):
122        for k in self._keys:
123            yield self[k]
124
125    def iteritems(self):
126        for k in self._keys:
127            yield k, self[k]
128
129    def keyrange(self, start=None, end=None, inclusive=False):
130        if start is not None:
131            start = self._left_ge(start)
132
133        if end is not None:
134            if inclusive:
135                end = self._right_le(end)
136            else:
137                end = self._right_lt(end)
138
139        return iter(self._keys[start:end+1])
140
141    def valuerange(self, *args, **kwargs):
142        for k in self.keyrange(*args, **kwargs):
143            yield self[k]
144
145    def itemrange(self, *args, **kwargs):
146        for k in self.keyrange(*args, **kwargs):
147            yield k, self[k]
148
149    def update(self, *args, **kwargs):
150        dict.update(self, *args, **kwargs)
151        self._del_keys()
152
153    def setdefault(self, key, _failobj=None):
154        try:
155            return self[key]
156        except KeyError:
157            self[key] = _failobj
158
159    def pop(self, key, *args):
160        try:
161            dict.pop(self, key)
162            self._del_keys()
163        except KeyError:
164            if not args:
165                raise
166            return args[0]
167
168    def popitem(self):
169        try:
170            key = self._keys[0]
171            self._del_keys()
172        except IndexError:
173            raise KeyError('popitem(): dictionary is empty')
174        else:
175            return key, dict.pop(self, key)
176
177    @classmethod
178    def fromkeys(cls, seq, value=None):
179        d = cls()
180        for key in seq:
181            d[key] = value
182        return d
183
184if __name__ == '__main__':
185    def display(d):
186        print(d)
187        print(d.keys())
188        print(list(d.iterkeys()))
189        print(d.values())
190        print(list(d.itervalues()))
191        print(d.items())
192        print(list(d.iteritems()))
193
194    d = SortedDict(x=24,e=5,j=4,b=2,z=26,d=4)
195    display(d)
196
197    print('popitem', d.popitem())
198    display(d)
199
200    print('pop j')
201    d.pop('j')
202    display(d)
203
204    d.setdefault('a', 1)
205    d.setdefault('g', 7)
206    d.setdefault('_')
207    display(d)
208
209    d.update({'b' : 2, 'h' : 8})
210    display(d)
211
212    del d['x']
213    display(d)
214    d['y'] = 26
215    display(d)
216
217    print(`d`)
218
219    print(d.copy())
220
221    for k,v in d.itemrange('d', 'z', inclusive=True):
222        print(k, v)
223