| 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 None: |
| 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.fd is None: |
| 180 | self.item = self._decode(self.cur.last()) |
| 181 | else: |
| 182 | k, v = self._decode(self.cur.set_range(self.typ.encode(self.ld))) |
| 183 | if self.fi: |
| 184 | while self.typ.compare(k, self.fd) == 0: |
| 185 | k, v = self._decode(self.cur.next()) |
| 186 | k, v = self._decode(self.cur.prev()) |
| 187 | else: |
| 188 | while self.typ.compare(k, self.fd) >= 0: |
| 189 | k, v = self._decode(self.cur.prev()) |
| 190 | self.item = k, v |
| 191 | except notfound: |
| 192 | self.item = StopIteration |
| 193 | |
| 194 | @dloopfun |
| 195 | def next(self): |
| 196 | try: |
| 197 | k, v = self.item = self._decode(self.cur.next()) |
| 198 | if ((self.li and self.typ.compare(k, self.ld) > 0) or |
| 199 | (not self.li and self.typ.compare(k, self.ld) >= 0)): |
| 200 | self.item = StopIteration |
| 201 | except notfound: |
| 202 | self.item = StopIteration |
| 203 | |
| 204 | @dloopfun |
| 205 | def prev(self): |
| 206 | try: |
| 207 | self.item = self._decode(self.cur.prev()) |
| 208 | if ((self.fi and self.typ.compare(k, self.fd) < 0) or |
| 209 | (not self.fi and self.typ.compare(k, self.fd) <= 0)): |
| 210 | self.item = StopIteration |
| 211 | except notfound: |
| 212 | self.item = StopIteration |
| 213 | |
| 214 | def __next__(self): |
| 215 | if self.item is None: |
| 216 | if not self.rev: |
| 217 | self.next() |
| 218 | else: |
| 219 | self.prev() |
| 220 | if self.item is StopIteration: |
| 221 | raise StopIteration() |
| 222 | ret, self.item = self.item, None |
| 223 | return ret |
| 224 | |
| 225 | def skip(self, n=1): |
| 226 | try: |
| 227 | for i in range(n): |
| 228 | next(self) |
| 229 | except StopIteration: |
| 230 | return |
| 231 | |
| 232 | def get(self, *, match=missing, ge=missing, gt=missing, lt=missing, le=missing, all=False, reverse=False): |
| 233 | if all: |
| 234 | cur = self.cursor(self, None, True, None, True, reverse) |
| 235 | elif match is not missing: |
| 236 | cur = self.cursor(self, match, True, match, True, reverse) |
| 237 | elif ge is not missing or gt is not missing or lt is not missing or le is not missing: |
| 238 | if ge is not missing: |
| 239 | fd, fi = ge, True |
| 240 | elif gt is not missing: |
| 241 | fd, fi = gt, False |
| 242 | else: |
| 243 | fd, fi = None, True |
| 244 | if le is not missing: |
| 245 | ld, li = le, True |
| 246 | elif lt is not missing: |
| 247 | ld, li = lt, False |
| 248 | else: |
| 249 | ld, li = None, True |
| 250 | cur = self.cursor(self, fd, fi, ld, li, reverse) |
| 251 | else: |
| 252 | raise NameError("invalid get() specification") |
| 253 | done = False |
| 254 | try: |
| 255 | if not reverse: |
| 256 | cur.first() |
| 257 | else: |
| 258 | cur.last() |
| 259 | done = True |
| 260 | return cur |
| 261 | finally: |
| 262 | if not done: |
| 263 | cur.close() |
| 264 | |
| 265 | @txnfun(lambda self: self.db.env.env) |
| 266 | def put(self, key, id, *, tx): |
| 267 | obid = struct.pack(">Q", id) |
| 268 | if not self.db.ob.has_key(obid, txn=tx.tx): |
| 269 | raise ValueError("no such object in database: " + str(id)) |
| 270 | try: |
| 271 | self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA) |
| 272 | except bd.DBKeyExistError: |
| 273 | return False |
| 274 | return True |
| 275 | |
| 276 | @txnfun(lambda self: self.db.env.env) |
| 277 | def remove(self, key, id, *, tx): |
| 278 | obid = struct.pack(">Q", id) |
| 279 | if not self.db.ob.has_key(obid, txn=tx.tx): |
| 280 | raise ValueError("no such object in database: " + str(id)) |
| 281 | cur = self.bk.cursor(txn=tx.tx) |
| 282 | try: |
| 283 | try: |
| 284 | cur.get_both(self.typ.encode(key), obid) |
| 285 | except notfound: |
| 286 | return False |
| 287 | cur.delete() |
| 288 | finally: |
| 289 | cur.close() |
| 290 | return True |