[didex.git] / didex / index.py

import struct, contextlib, math
from . import db, lib
from .db import bd, txnfun

__all__ = ["maybe", "t_int", "t_uint", "t_float", "t_str", "ordered"]

deadlock = bd.DBLockDeadlockError
notfound = bd.DBNotFoundError

class simpletype(object):
    def __init__(self, encode, decode):
        self.enc = encode
        self.dec = decode

    def encode(self, ob):
        return self.enc(ob)
    def decode(self, dat):
        return self.dec(dat)
    def compare(self, a, b):
        if a < b:
            return -1
        elif a > b:
            return 1
        else:
            return 0

    @classmethod
    def struct(cls, fmt):
        return cls(lambda ob: struct.pack(fmt, ob),
                   lambda dat: struct.unpack(fmt, dat)[0])

class maybe(object):
    def __init__(self, bk):
        self.bk = bk

    def encode(self, ob):
        if ob is None: return b""
        return b"\0" + self.bk.encode(ob)
    def decode(self, dat):
        if dat == b"": return None
        return self.bk.dec(dat[1:])
    def compare(self, a, b):
        if a is b is None:
            return 0
        elif a is None:
            return -1
        elif b is None:
            return 1
        else:
            return self.bk.compare(a[1:], b[1:])

class compound(object):
    def __init__(self, *parts):
        self.parts = parts

    def encode(self, obs):
        if len(obs) != len(self.parts):
            raise ValueError("invalid length of compound data: " + str(len(obs)) + ", rather than " + len(self.parts))
        buf = bytearray()
        for ob, part in zip(obs, self.parts):
            dat = part.encode(ob)
            if len(dat) < 128:
                buf.append(0x80 | len(dat))
                buf.extend(dat)
            else:
                buf.extend(struct.pack(">i", len(dat)))
                buf.extend(dat)
        return bytes(buf)
    def decode(self, dat):
        ret = []
        off = 0
        for part in self.parts:
            if dat[off] & 0x80:
                ln = dat[off] & 0x7f
                off += 1
            else:
                ln = struct.unpack(">i", dat[off:off + 4])[0]
                off += 4
            ret.append(part.decode(dat[off:off + len]))
            off += len
        return tuple(ret)
    def compare(self, al, bl):
        if (len(al) != len(self.parts)) or (len(bl) != len(self.parts)):
            raise ValueError("invalid length of compound data: " + str(len(al)) + ", " + str(len(bl)) + ", rather than " + len(self.parts))
        for a, b, part in zip(al, bl, self.parts):
            c = part.compare(a, b)
            if c != 0:
                return c
        return 0

def floatcmp(a, b):
    if math.isnan(a) and math.isnan(b):
        return 0
    elif math.isnan(a):
        return -1
    elif math.isnan(b):
        return 1
    elif a < b:
        return -1
    elif a > b:
        return 1
    else:
        return 0

t_int = simpletype.struct(">q")
t_uint = simpletype.struct(">Q")
t_float = simpletype.struct(">d")
t_float.compare = floatcmp
t_str = simpletype((lambda ob: ob.encode("utf-8")), (lambda dat: dat.decode("utf-8")))

class index(object):
    def __init__(self, db, name, datatype):
        self.db = db
        self.nm = name
        self.typ = datatype

missing = object()

class ordered(index, lib.closable):
    def __init__(self, db, name, datatype, create=True):
        super().__init__(db, name, datatype)
        fl = bd.DB_THREAD | bd.DB_AUTO_COMMIT
        if create: fl |= bd.DB_CREATE
        def initdb(db):
            def compare(a, b):
                if a == b == "": return 0
                return self.typ.compare(self.typ.decode(a), self.typ.decode(b))
            db.set_flags(bd.DB_DUPSORT)
            db.set_bt_compare(compare)
        self.bk = db._opendb("i-" + name, bd.DB_BTREE, fl, initdb)
        self.bk.set_get_returns_none(False)

    def close(self):
        self.bk.close()

    class cursor(lib.closable):
        def __init__(self, idx, cur, item, stop):
            self.idx = idx
            self.cur = cur
            self.item = item
            self.stop = stop

        def close(self):
            if self.cur is not None:
                self.cur.close()

        def __iter__(self):
            return self

        def peek(self):
            if self.item is None:
                raise StopIteration()
            rk, rv = self.item
            rk = self.idx.typ.decode(rk)
            rv = struct.unpack(">Q", rv)[0]
            if self.stop(rk):
                self.item = None
                raise StopIteration()
            return rk, rv

        def __next__(self):
            rk, rv = self.peek()
            try:
                while True:
                    try:
                        self.item = self.cur.next()
                        break
                    except deadlock:
                        continue
            except notfound:
                self.item = None
            return rk, rv

        def skip(self, n=1):
            try:
                for i in range(n):
                    next(self)
            except StopIteration:
                return

    def get(self, *, match=missing, ge=missing, gt=missing, lt=missing, le=missing, all=False):
        while True:
            try:
                cur = self.bk.cursor()
                done = False
                try:
                    if match is not missing:
                        try:
                            k, v = cur.set(self.typ.encode(match))
                        except notfound:
                            return self.cursor(None, None, None, None)
                        else:
                            done = True
                            return self.cursor(self, cur, (k, v), lambda o: (self.typ.compare(o, match) != 0))
                    elif all:
                        try:
                            k, v = cur.first()
                        except notfound:
                            return self.cursor(None, None, None, None)
                        else:
                            done = True
                            return self.cursor(self, cur, (k, v), lambda o: False)
                    elif ge is not missing or gt is not missing or lt is not missing or le is not missing:
                        skip = False
                        try:
                            if ge is not missing:
                                k, v = cur.set_range(self.typ.encode(ge))
                            elif gt is not missing:
                                k, v = cur.set_range(self.typ.encode(gt))
                                skip = True
                            else:
                                k, v = cur.first()
                        except notfound:
                            return self.cursor(None, None, None, None)
                        if lt is not missing:
                            stop = lambda o: self.typ.compare(o, lt) >= 0
                        elif le is not missing:
                            stop = lambda o: self.typ.compare(o, le) > 0
                        else:
                            stop = lambda o: False
                        ret = self.cursor(self, cur, (k, v), stop)
                        if skip:
                            try:
                                while self.typ.compare(ret.peek()[0], gt) == 0:
                                    next(ret)
                            except StopIteration:
                                pass
                        done = True
                        return ret
                    else:
                        raise NameError("invalid get() specification")
                finally:
                    if not done:
                        cur.close()
            except deadlock:
                continue

    @txnfun(lambda self: self.db.env.env)
    def put(self, key, id, *, tx):
        obid = struct.pack(">Q", id)
        if not self.db.ob.has_key(obid, txn=tx.tx):
            raise ValueError("no such object in database: " + str(id))
        try:
            self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA)
        except bd.DBKeyExistError:
            return False
        return True

    @txnfun(lambda self: self.db.env.env)
    def remove(self, key, id, *, tx):
        obid = struct.pack(">Q", id)
        if not self.db.ob.has_key(obid, txn=tx.tx):
            raise ValueError("no such object in database: " + str(id))
        cur = self.bk.cursor(txn=tx.tx)
        try:
            try:
                cur.get_both(self.typ.encode(key), obid)
            except notfound:
                return False
            cur.delete()
        finally:
            cur.close()
        return True
Commit	Line	Data
abb94f83	1	import struct, contextlib, math
a95055e8	2	from . import db, lib
8950191c	3	from .db import bd, txnfun
a95055e8	4
cbf73d3a FT	5	__all__ = ["maybe", "t_int", "t_uint", "t_float", "t_str", "ordered"]
cbf73d3a FT	6
a95055e8 FT	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
bd14729f FT	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
abb94f83 FT	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")))
a95055e8 FT	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):
eb274691	120	def __init__(self, db, name, datatype, create=True):
a95055e8	121	super().__init__(db, name, datatype)
a95055e8 FT	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, cur, item, stop):
	138	self.idx = idx
	139	self.cur = cur
	140	self.item = item
	141	self.stop = stop
	142
	143	def close(self):
	144	if self.cur is not None:
	145	self.cur.close()
	146
	147	def __iter__(self):
	148	return self
	149
	150	def peek(self):
	151	if self.item is None:
	152	raise StopIteration()
	153	rk, rv = self.item
	154	rk = self.idx.typ.decode(rk)
	155	rv = struct.unpack(">Q", rv)[0]
	156	if self.stop(rk):
	157	self.item = None
	158	raise StopIteration()
	159	return rk, rv
	160
	161	def __next__(self):
	162	rk, rv = self.peek()
	163	try:
	164	while True:
	165	try:
	166	self.item = self.cur.next()
	167	break
	168	except deadlock:
	169	continue
	170	except notfound:
	171	self.item = None
	172	return rk, rv
	173
	174	def skip(self, n=1):
	175	try:
	176	for i in range(n):
	177	next(self)
	178	except StopIteration:
	179	return
	180
	181	def get(self, *, match=missing, ge=missing, gt=missing, lt=missing, le=missing, all=False):
	182	while True:
	183	try:
	184	cur = self.bk.cursor()
	185	done = False
186	try:
187	if match is not missing:
188	try:
189	k, v = cur.set(self.typ.encode(match))
190	except notfound:
191	return self.cursor(None, None, None, None)
192	else:
193	done = True
194	return self.cursor(self, cur, (k, v), lambda o: (self.typ.compare(o, match) != 0))
195	elif all:
196	try:
197	k, v = cur.first()
198	except notfound:
199	return self.cursor(None, None, None, None)
200	else:
201	done = True
202	return self.cursor(self, cur, (k, v), lambda o: False)
203	elif ge is not missing or gt is not missing or lt is not missing or le is not missing:
204	skip = False
205	try:
206	if ge is not missing:
207	k, v = cur.set_range(self.typ.encode(ge))
208	elif gt is not missing:
209	k, v = cur.set_range(self.typ.encode(gt))
210	skip = True
211	else:
212	k, v = cur.first()
213	except notfound:
214	return self.cursor(None, None, None, None)
215	if lt is not missing:
216	stop = lambda o: self.typ.compare(o, lt) >= 0
217	elif le is not missing:
218	stop = lambda o: self.typ.compare(o, le) > 0
219	else:
220	stop = lambda o: False
221	ret = self.cursor(self, cur, (k, v), stop)
222	if skip:
223	try:
224	while self.typ.compare(ret.peek()[0], gt) == 0:
225	next(ret)
226	except StopIteration:
227	pass
228	done = True
229	return ret
230	else:
231	raise NameError("invalid get() specification")
232	finally:
233	if not done:
234	cur.close()
235	except deadlock:
236	continue
237
8950191c FT	238	@txnfun(lambda self: self.db.env.env)
	239	def put(self, key, id, *, tx):
	240	obid = struct.pack(">Q", id)
	241	if not self.db.ob.has_key(obid, txn=tx.tx):
	242	raise ValueError("no such object in database: " + str(id))
	243	try:
	244	self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA)
	245	except bd.DBKeyExistError:
	246	return False
	247	return True
	248
	249	@txnfun(lambda self: self.db.env.env)
	250	def remove(self, key, id, *, tx):
	251	obid = struct.pack(">Q", id)
	252	if not self.db.ob.has_key(obid, txn=tx.tx):
	253	raise ValueError("no such object in database: " + str(id))
	254	cur = self.bk.cursor(txn=tx.tx)
	255	try:
a95055e8	256	try:
8950191c FT	257	cur.get_both(self.typ.encode(key), obid)
	258	except notfound:
	259	return False
	260	cur.delete()
	261	finally:
	262	cur.close()
	263	return True