951eb3f3a942420bb376d26b66ab5fd9160ebdab
[didex.git] / didex / index.py
1 import struct, contextlib, math
2 from . import db, lib
3 from .db import bd, txnfun, dloopfun
4
5 __all__ = ["maybe", "t_int", "t_uint", "t_float", "t_str", "ordered"]
6
7 deadlock = bd.DBLockDeadlockError
8 notfound = bd.DBNotFoundError
9
10 class simpletype(object):
11     def __init__(self, encode, decode):
12         self.enc = encode
13         self.dec = decode
14
15     def encode(self, ob):
16         return self.enc(ob)
17     def decode(self, dat):
18         return self.dec(dat)
19     def compare(self, a, b):
20         if a < b:
21             return -1
22         elif a > b:
23             return 1
24         else:
25             return 0
26
27     @classmethod
28     def struct(cls, fmt):
29         return cls(lambda ob: struct.pack(fmt, ob),
30                    lambda dat: struct.unpack(fmt, dat)[0])
31
32 class maybe(object):
33     def __init__(self, bk):
34         self.bk = bk
35
36     def encode(self, ob):
37         if ob is None: return b""
38         return b"\0" + self.bk.encode(ob)
39     def decode(self, dat):
40         if dat == b"": return None
41         return self.bk.dec(dat[1:])
42     def compare(self, a, b):
43         if a is b is None:
44             return 0
45         elif a is None:
46             return -1
47         elif b is None:
48             return 1
49         else:
50             return self.bk.compare(a[1:], b[1:])
51
52 class compound(object):
53     def __init__(self, *parts):
54         self.parts = parts
55
56     def encode(self, obs):
57         if len(obs) != len(self.parts):
58             raise ValueError("invalid length of compound data: " + str(len(obs)) + ", rather than " + len(self.parts))
59         buf = bytearray()
60         for ob, part in zip(obs, self.parts):
61             dat = part.encode(ob)
62             if len(dat) < 128:
63                 buf.append(0x80 | len(dat))
64                 buf.extend(dat)
65             else:
66                 buf.extend(struct.pack(">i", len(dat)))
67                 buf.extend(dat)
68         return bytes(buf)
69     def decode(self, dat):
70         ret = []
71         off = 0
72         for part in self.parts:
73             if dat[off] & 0x80:
74                 ln = dat[off] & 0x7f
75                 off += 1
76             else:
77                 ln = struct.unpack(">i", dat[off:off + 4])[0]
78                 off += 4
79             ret.append(part.decode(dat[off:off + len]))
80             off += len
81         return tuple(ret)
82     def compare(self, al, bl):
83         if (len(al) != len(self.parts)) or (len(bl) != len(self.parts)):
84             raise ValueError("invalid length of compound data: " + str(len(al)) + ", " + str(len(bl)) + ", rather than " + len(self.parts))
85         for a, b, part in zip(al, bl, self.parts):
86             c = part.compare(a, b)
87             if c != 0:
88                 return c
89         return 0
90
91 def floatcmp(a, b):
92     if math.isnan(a) and math.isnan(b):
93         return 0
94     elif math.isnan(a):
95         return -1
96     elif math.isnan(b):
97         return 1
98     elif a < b:
99         return -1
100     elif a > b:
101         return 1
102     else:
103         return 0
104
105 t_int = simpletype.struct(">q")
106 t_uint = simpletype.struct(">Q")
107 t_float = simpletype.struct(">d")
108 t_float.compare = floatcmp
109 t_str = simpletype((lambda ob: ob.encode("utf-8")), (lambda dat: dat.decode("utf-8")))
110
111 class index(object):
112     def __init__(self, db, name, datatype):
113         self.db = db
114         self.nm = name
115         self.typ = datatype
116
117 missing = object()
118
119 class ordered(index, lib.closable):
120     def __init__(self, db, name, datatype, create=True):
121         super().__init__(db, name, datatype)
122         fl = bd.DB_THREAD | bd.DB_AUTO_COMMIT
123         if create: fl |= bd.DB_CREATE
124         def initdb(db):
125             def compare(a, b):
126                 if a == b == "": return 0
127                 return self.typ.compare(self.typ.decode(a), self.typ.decode(b))
128             db.set_flags(bd.DB_DUPSORT)
129             db.set_bt_compare(compare)
130         self.bk = db._opendb("i-" + name, bd.DB_BTREE, fl, initdb)
131         self.bk.set_get_returns_none(False)
132
133     def close(self):
134         self.bk.close()
135
136     class cursor(lib.closable):
137         def __init__(self, idx, fd, fi, ld, li, reverse):
138             self.idx = idx
139             self.typ = idx.typ
140             self.cur = self.idx.bk.cursor()
141             self.item = None
142             self.fd = fd
143             self.fi = fi
144             self.ld = ld
145             self.li = li
146             self.rev = reverse
147
148         def close(self):
149             if self.cur is not None:
150                 self.cur.close()
151                 self.cur = None
152
153         def __iter__(self):
154             return self
155
156         def _decode(self, d):
157             k, v = d
158             k = self.typ.decode(k)
159             v = struct.unpack(">Q", v)[0]
160             return k, v
161
162         @dloopfun
163         def first(self):
164             try:
165                 if self.fd is missing:
166                     self.item = self._decode(self.cur.first())
167                 else:
168                     k, v = self._decode(self.cur.set_range(self.typ.encode(self.fd)))
169                     if not self.fi:
170                         while self.typ.compare(k, self.fd) == 0:
171                             k, v = self._decode(self.cur.next())
172                     self.item = k, v
173             except notfound:
174                 self.item = StopIteration
175
176         @dloopfun
177         def last(self):
178             try:
179                 if self.ld is missing:
180                     self.item = self._decode(self.cur.last())
181                 else:
182                     try:
183                         k, v = self._decode(self.cur.set_range(self.typ.encode(self.ld)))
184                     except notfound:
185                         k, v = self._decode(self.cur.last())
186                     if self.li:
187                         while self.typ.compare(k, self.ld) == 0:
188                             k, v = self._decode(self.cur.next())
189                         while self.typ.compare(k, self.ld) > 0:
190                             k, v = self._decode(self.cur.prev())
191                     else:
192                         while self.typ.compare(k, self.ld) >= 0:
193                             k, v = self._decode(self.cur.prev())
194                     self.item = k, v
195             except notfound:
196                 self.item = StopIteration
197
198         @dloopfun
199         def next(self):
200             try:
201                 k, v = self.item = self._decode(self.cur.next())
202                 if (self.ld is not missing and
203                     ((self.li and self.typ.compare(k, self.ld) > 0) or
204                      (not self.li and self.typ.compare(k, self.ld) >= 0))):
205                     self.item = StopIteration
206             except notfound:
207                 self.item = StopIteration
208
209         @dloopfun
210         def prev(self):
211             try:
212                 self.item = self._decode(self.cur.prev())
213                 if (self.fd is not missing and
214                     ((self.fi and self.typ.compare(k, self.fd) < 0) or
215                      (not self.fi and self.typ.compare(k, self.fd) <= 0))):
216                     self.item = StopIteration
217             except notfound:
218                 self.item = StopIteration
219
220         def __next__(self):
221             if self.item is None:
222                 if not self.rev:
223                     self.next()
224                 else:
225                     self.prev()
226             if self.item is StopIteration:
227                 raise StopIteration()
228             ret, self.item = self.item, None
229             return ret
230
231         def skip(self, n=1):
232             try:
233                 for i in range(n):
234                     next(self)
235             except StopIteration:
236                 return
237
238     def get(self, *, match=missing, ge=missing, gt=missing, lt=missing, le=missing, all=False, reverse=False):
239         if all:
240             cur = self.cursor(self, missing, True, missing, True, reverse)
241         elif match is not missing:
242             cur = self.cursor(self, match, True, match, True, reverse)
243         elif ge is not missing or gt is not missing or lt is not missing or le is not missing:
244             if ge is not missing:
245                 fd, fi = ge, True
246             elif gt is not missing:
247                 fd, fi = gt, False
248             else:
249                 fd, fi = missing, True
250             if le is not missing:
251                 ld, li = le, True
252             elif lt is not missing:
253                 ld, li = lt, False
254             else:
255                 ld, li = missing, True
256             cur = self.cursor(self, fd, fi, ld, li, reverse)
257         else:
258             raise NameError("invalid get() specification")
259         done = False
260         try:
261             if not reverse:
262                 cur.first()
263             else:
264                 cur.last()
265             done = True
266             return cur
267         finally:
268             if not done:
269                 cur.close()
270
271     @txnfun(lambda self: self.db.env.env)
272     def put(self, key, id, *, tx):
273         obid = struct.pack(">Q", id)
274         if not self.db.ob.has_key(obid, txn=tx.tx):
275             raise ValueError("no such object in database: " + str(id))
276         try:
277             self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA)
278         except bd.DBKeyExistError:
279             return False
280         return True
281
282     @txnfun(lambda self: self.db.env.env)
283     def remove(self, key, id, *, tx):
284         obid = struct.pack(">Q", id)
285         if not self.db.ob.has_key(obid, txn=tx.tx):
286             raise ValueError("no such object in database: " + str(id))
287         cur = self.bk.cursor(txn=tx.tx)
288         try:
289             try:
290                 cur.get_both(self.typ.encode(key), obid)
291             except notfound:
292                 return False
293             cur.delete()
294         finally:
295             cur.close()
296         return True