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
28from __future__ import absolute_import
29
30from bisect import bisect_left, bisect_right
31
32class SortedDict(dict):
33    def _get_sorted(self):
34        return getattr(self, '_sorted', sorted)
35    def _set_sorted(self, val):
36        self._sorted = val
37        self._del_keys()
38    sorted = property(_get_sorted, _set_sorted)
39
40    @property
41    def _keys(self):
42        try:
43            return self._sorted_keys
44        except AttributeError:
45            _sorted_keys = self.sorted(dict.keys(self))
46            self._sorted_keys = _sorted_keys
47            return _sorted_keys
48
49    def _left_eq(self, key):
50        index = self._left_ge(self, key)
51        if self._keys[index] != key:
52            raise KeyError(key)
53        return index
54
55    def _right_eq(self, key):
56        index = self._right_le(self, key)
57        if self._keys[index] != key:
58            raise KeyError(key)
59        return index
60
61    def _right_lt(self, key):
62        index = bisect_left(self._keys, key)
63        if index:
64            return index - 1
65        raise KeyError(key)
66
67    def _right_le(self, key):
68        index = bisect_right(self._keys, key)
69        if index:
70            return index - 1
71        raise KeyError(key)
72
73    def _left_gt(self, key):
74        index = bisect_right(self._keys, key)
75        if index != len(self._keys):
76            return index
77        raise KeyError(key)
78
79    def _left_ge(self, key):
80        index = bisect_left(self._keys, key)
81        if index != len(self._keys):
82            return index
83        raise KeyError(key)
84
85    def _del_keys(self):
86        try:
87            del self._sorted_keys
88        except AttributeError:
89            pass
90
91    def __repr__(self):
92        return 'SortedDict({%s})' % ', '.join('%r: %r' % item
93                                              for item in self.items())
94    def __setitem__(self, key, item):
95        dict.__setitem__(self, key, item)
96        self._del_keys()
97
98    def __delitem__(self, key):
99        dict.__delitem__(self, key)
100        self._del_keys()
101
102    def clear(self):
103        self.data.clear()
104        self._del_keys()
105
106    def copy(self):
107        t = type(self)
108        return t(self)
109
110    def keys(self):
111        return self._keys
112
113    def values(self):
114        for k in self._keys:
115            yield self[k]
116
117    def items(self):
118        for k in self._keys:
119            yield k, self[k]
120
121    def keyrange(self, start=None, end=None, inclusive=False):
122        if start is not None:
123            start = self._left_ge(start)
124
125        if end is not None:
126            if inclusive:
127                end = self._right_le(end)
128            else:
129                end = self._right_lt(end)
130
131        return iter(self._keys[start:end+1])
132
133    def valuerange(self, *args, **kwargs):
134        for k in self.keyrange(*args, **kwargs):
135            yield self[k]
136
137    def itemrange(self, *args, **kwargs):
138        for k in self.keyrange(*args, **kwargs):
139            yield k, self[k]
140
141    def update(self, *args, **kwargs):
142        dict.update(self, *args, **kwargs)
143        self._del_keys()
144
145    def setdefault(self, key, _failobj=None):
146        try:
147            return self[key]
148        except KeyError:
149            self[key] = _failobj
150
151    def pop(self, key, *args):
152        try:
153            dict.pop(self, key)
154            self._del_keys()
155        except KeyError:
156            if not args:
157                raise
158            return args[0]
159
160    def popitem(self):
161        try:
162            key = self._keys[0]
163            self._del_keys()
164        except IndexError:
165            raise KeyError('popitem(): dictionary is empty')
166        else:
167            return key, dict.pop(self, key)
168
169    @classmethod
170    def fromkeys(cls, seq, value=None):
171        d = cls()
172        for key in seq:
173            d[key] = value
174        return d
175
176if __name__ == '__main__':
177    def display(d):
178        print(d)
179        print(list(d.keys()))
180        print(list(d.keys()))
181        print(list(d.values()))
182        print(list(d.values()))
183        print(list(d.items()))
184        print(list(d.items()))
185
186    d = SortedDict(x=24,e=5,j=4,b=2,z=26,d=4)
187    display(d)
188
189    print('popitem', d.popitem())
190    display(d)
191
192    print('pop j')
193    d.pop('j')
194    display(d)
195
196    d.setdefault('a', 1)
197    d.setdefault('g', 7)
198    d.setdefault('_')
199    display(d)
200
201    d.update({'b' : 2, 'h' : 8})
202    display(d)
203
204    del d['x']
205    display(d)
206    d['y'] = 26
207    display(d)
208
209    print(repr(d))
210
211    print(d.copy())
212
213    for k,v in d.itemrange('d', 'z', inclusive=True):
214        print(k, v)
215