summaryrefslogtreecommitdiff
path: root/blist/_sorteddict.py
blob: fcdf7e41cd6f59299b375cfb1bca9844085b4c28 (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
from blist._sortedlist import sortedset, ReprRecursion
import collections, sys
from blist._blist import blist

class missingdict(dict):
    def __missing__(self, key):
        return self._missing(key)

class KeysView(collections.KeysView, collections.Sequence):
    def __getitem__(self, index):
        return self._mapping._sortedkeys[index]
    def __reversed__(self):
        return reversed(self._mapping._sortedkeys)
    def index(self, key):
        return self._mapping._sortedkeys.index(key)
    def count(self, key):
        return 1 if key in self else 0
    def _from_iterable(self, it):
        return sortedset(it, key=self._mapping._sortedkeys._key)
    def bisect_left(self, key):
        return self._mapping._sortedkeys.bisect_left(key)
    def bisect_right(self, key):
        return self._mapping._sortedkeys.bisect_right(key)
    bisect = bisect_right

class ItemsView(collections.ItemsView, collections.Sequence):
    def __getitem__(self, index):
        if isinstance(index, slice):
            keys = self._mapping._sortedkeys[index]
            return self._from_iterable((key, self._mapping[key])
                                       for key in keys)
        key = self._mapping._sortedkeys[index]
        return (key, self._mapping[key])
    def index(self, item):
        key, value = item
        i = self._mapping._sortedkeys.index(key)
        if self._mapping[key] == value:
            return i
        raise ValueError
    def count(self, item):
        return 1 if item in self else 0
    def _from_iterable(self, it):
      keyfunc = self._mapping._sortedkeys._key
      if keyfunc is None:
        return sortedset(it)
      else:
        return sortedset(it, key=lambda item: keyfunc(item[0]))

class ValuesView(collections.ValuesView, collections.Sequence):
    def __getitem__(self, index):
        if isinstance(index, slice):
            keys = self._mapping._sortedkeys[index]
            return [self._mapping[key] for key in keys]
        key = self._mapping._sortedkeys[index]
        return self._mapping[key]

class sorteddict(collections.MutableMapping):
    def __init__(self, *args, **kw):
        if hasattr(self, '__missing__'):
            self._map = missingdict()
            self._map._missing = self.__missing__
        else:
            self._map = dict()
        key = None
        if len(args) > 0:
            if hasattr(args[0], '__call__'):
                key = args[0]
                args = args[1:]
            elif len(args) > 1:
                raise TypeError("'%s' object is not callable" %
                                args[0].__class__.__name__)
        if len(args) > 1:
            raise TypeError('sorteddict expected at most 2 arguments, got %d'
                            % len(args))
        if len(args) == 1 and isinstance(args[0], sorteddict) and key is None:
            key = args[0]._sortedkeys._key
        self._sortedkeys = sortedset(key=key)
        self.update(*args, **kw)

    if sys.version_info[0] < 3:
        def keys(self):
            return self._sortedkeys.copy()
        def items(self):
            return blist((key, self[key]) for key in self)
        def values(self):
            return blist(self[key] for key in self)
        def viewkeys(self):
            return KeysView(self)
        def viewitems(self):
            return ItemsView(self)
        def viewvalues(self):
            return ValuesView(self)
    else:
        def keys(self):
            return KeysView(self)
        def items(self):
            return ItemsView(self)
        def values(self):
            return ValuesView(self)

    def __setitem__(self, key, value):
        try:
            if key not in self._map:
                self._sortedkeys.add(key)
            self._map[key] = value
        except:
            if key not in self._map:
                self._sortedkeys.discard(key)
            raise

    def __delitem__(self, key):
        self._sortedkeys.discard(key)
        del self._map[key]

    def __getitem__(self, key):
        return self._map[key]

    def __iter__(self):
        return iter(self._sortedkeys)

    def __len__(self):
        return len(self._sortedkeys)

    def copy(self):
        return sorteddict(self)

    @classmethod
    def fromkeys(cls, keys, value=None, key=None):
        if key is not None:
            rv = cls(key)
        else:
            rv = cls()
        for key in keys:
            rv[key] = value
        return rv

    def __repr__(self):
        with ReprRecursion(self) as r:
            if r:
              return 'sorteddict({...})'
            return ('sorteddict({%s})' %
                    ', '.join('%r: %r' % (k, self._map[k]) for k in self))

    def __eq__(self, other):
        if not isinstance(other, sorteddict):
            return False
        return self._map == other._map