[didex.git] / didex / index.py

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

__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

    small = object()
    large = object()
    def minim(self, *parts):
        return parts + tuple([self.small] * (len(self.parts) - len(parts)))
    def maxim(self, *parts):
        return parts + tuple([self.large] * (len(self.parts) - len(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):
            if ob is self.small:
                buf.append(0x01)
            elif ob is self.large:
                buf.append(0x02)
            else:
                dat = part.encode(ob)
                if len(dat) < 128:
                    buf.append(0x80 | len(dat))
                    buf.extend(dat)
                else:
                    buf.extend(struct.pack(">BI", 0, len(dat)))
                    buf.extend(dat)
        return bytes(buf)
    def decode(self, dat):
        ret = []
        off = 0
        for part in self.parts:
            fl = dat[off]
            off += 1
            if fl & 0x80:
                ln = fl & 0x7f
            elif fl == 0x01:
                ret.append(self.small)
                continue
            elif fl == 0x02:
                ret.append(self.large)
                continue
            else:
                ln = struct.unpack(">I", dat[off:off + 4])[0]
                off += 4
            ret.append(part.decode(dat[off:off + ln]))
            off += ln
        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):
            if a in (self.small, self.large) or b in (self.small, self.large):
                if a is b:
                    return 0
                if a is self.small:
                    return -1
                elif b is self.small:
                    return 1
                elif a is self.large:
                    return 1
                elif b is self.large:
                    return -1
            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, fd, fi, ld, li, reverse):
            self.idx = idx
            self.typ = idx.typ
            self.cur = self.idx.bk.cursor()
            self.item = None
            self.fd = fd
            self.fi = fi
            self.ld = ld
            self.li = li
            self.rev = reverse

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

        def __iter__(self):
            return self

        def _decode(self, d):
            k, v = d
            k = self.typ.decode(k)
            v = struct.unpack(">Q", v)[0]
            return k, v

        @dloopfun
        def first(self):
            try:
                if self.fd is missing:
                    self.item = self._decode(self.cur.first())
                else:
                    k, v = self._decode(self.cur.set_range(self.typ.encode(self.fd)))
                    if not self.fi:
                        while self.typ.compare(k, self.fd) == 0:
                            k, v = self._decode(self.cur.next())
                    self.item = k, v
            except notfound:
                self.item = StopIteration

        @dloopfun
        def last(self):
            try:
                if self.ld is missing:
                    self.item = self._decode(self.cur.last())
                else:
                    try:
                        k, v = self._decode(self.cur.set_range(self.typ.encode(self.ld)))
                    except notfound:
                        k, v = self._decode(self.cur.last())
                    if self.li:
                        while self.typ.compare(k, self.ld) == 0:
                            k, v = self._decode(self.cur.next())
                        while self.typ.compare(k, self.ld) > 0:
                            k, v = self._decode(self.cur.prev())
                    else:
                        while self.typ.compare(k, self.ld) >= 0:
                            k, v = self._decode(self.cur.prev())
                    self.item = k, v
            except notfound:
                self.item = StopIteration

        @dloopfun
        def next(self):
            try:
                k, v = self.item = self._decode(self.cur.next())
                if (self.ld is not missing and
                    ((self.li and self.typ.compare(k, self.ld) > 0) or
                     (not self.li and self.typ.compare(k, self.ld) >= 0))):
                    self.item = StopIteration
            except notfound:
                self.item = StopIteration

        @dloopfun
        def prev(self):
            try:
                self.item = self._decode(self.cur.prev())
                if (self.fd is not missing and
                    ((self.fi and self.typ.compare(k, self.fd) < 0) or
                     (not self.fi and self.typ.compare(k, self.fd) <= 0))):
                    self.item = StopIteration
            except notfound:
                self.item = StopIteration

        def __next__(self):
            if self.item is None:
                if not self.rev:
                    self.next()
                else:
                    self.prev()
            if self.item is StopIteration:
                raise StopIteration()
            ret, self.item = self.item, None
            return ret

        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, reverse=False):
        if all:
            cur = self.cursor(self, missing, True, missing, True, reverse)
        elif match is not missing:
            cur = self.cursor(self, match, True, match, True, reverse)
        elif ge is not missing or gt is not missing or lt is not missing or le is not missing:
            if ge is not missing:
                fd, fi = ge, True
            elif gt is not missing:
                fd, fi = gt, False
            else:
                fd, fi = missing, True
            if le is not missing:
                ld, li = le, True
            elif lt is not missing:
                ld, li = lt, False
            else:
                ld, li = missing, True
            cur = self.cursor(self, fd, fi, ld, li, reverse)
        else:
            raise NameError("invalid get() specification")
        done = False
        try:
            if not reverse:
                cur.first()
            else:
                cur.last()
            done = True
            return cur
        finally:
            if not done:
                cur.close()

    @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
6efe4e23	3	from .db import bd, txnfun, dloopfun
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
177fbee6 FT	56	small = object()
	57	large = object()
	58	def minim(self, *parts):
	59	return parts + tuple([self.small] * (len(self.parts) - len(parts)))
	60	def maxim(self, *parts):
	61	return parts + tuple([self.large] * (len(self.parts) - len(parts)))
	62
bd14729f FT	63	def encode(self, obs):
	64	if len(obs) != len(self.parts):
	65	raise ValueError("invalid length of compound data: " + str(len(obs)) + ", rather than " + len(self.parts))
	66	buf = bytearray()
	67	for ob, part in zip(obs, self.parts):
177fbee6 FT	68	if ob is self.small:
	69	buf.append(0x01)
	70	elif ob is self.large:
	71	buf.append(0x02)
bd14729f	72	else:
177fbee6 FT	73	dat = part.encode(ob)
	74	if len(dat) < 128:
	75	buf.append(0x80 \| len(dat))
	76	buf.extend(dat)
	77	else:
	78	buf.extend(struct.pack(">BI", 0, len(dat)))
	79	buf.extend(dat)
bd14729f FT	80	return bytes(buf)
	81	def decode(self, dat):
	82	ret = []
	83	off = 0
	84	for part in self.parts:
177fbee6 FT	85	fl = dat[off]
	86	off += 1
	87	if fl & 0x80:
	88	ln = fl & 0x7f
	89	elif fl == 0x01:
	90	ret.append(self.small)
	91	continue
	92	elif fl == 0x02:
	93	ret.append(self.large)
	94	continue
bd14729f	95	else:
177fbee6	96	ln = struct.unpack(">I", dat[off:off + 4])[0]
bd14729f	97	off += 4
177fbee6 FT	98	ret.append(part.decode(dat[off:off + ln]))
177fbee6 FT	99	off += ln
bd14729f FT	100	return tuple(ret)
	101	def compare(self, al, bl):
	102	if (len(al) != len(self.parts)) or (len(bl) != len(self.parts)):
	103	raise ValueError("invalid length of compound data: " + str(len(al)) + ", " + str(len(bl)) + ", rather than " + len(self.parts))
	104	for a, b, part in zip(al, bl, self.parts):
177fbee6 FT	105	if a in (self.small, self.large) or b in (self.small, self.large):
	106	if a is b:
	107	return 0
	108	if a is self.small:
	109	return -1
	110	elif b is self.small:
	111	return 1
	112	elif a is self.large:
	113	return 1
	114	elif b is self.large:
	115	return -1
bd14729f FT	116	c = part.compare(a, b)
	117	if c != 0:
	118	return c
	119	return 0
	120
abb94f83 FT	121	def floatcmp(a, b):
	122	if math.isnan(a) and math.isnan(b):
	123	return 0
	124	elif math.isnan(a):
	125	return -1
	126	elif math.isnan(b):
	127	return 1
	128	elif a < b:
	129	return -1
	130	elif a > b:
	131	return 1
	132	else:
	133	return 0
	134
	135	t_int = simpletype.struct(">q")
	136	t_uint = simpletype.struct(">Q")
	137	t_float = simpletype.struct(">d")
	138	t_float.compare = floatcmp
	139	t_str = simpletype((lambda ob: ob.encode("utf-8")), (lambda dat: dat.decode("utf-8")))
a95055e8 FT	140
	141	class index(object):
	142	def __init__(self, db, name, datatype):
	143	self.db = db
	144	self.nm = name
	145	self.typ = datatype
	146
	147	missing = object()
	148
	149	class ordered(index, lib.closable):
eb274691	150	def __init__(self, db, name, datatype, create=True):
a95055e8	151	super().__init__(db, name, datatype)
a95055e8 FT	152	fl = bd.DB_THREAD \| bd.DB_AUTO_COMMIT
	153	if create: fl \|= bd.DB_CREATE
	154	def initdb(db):
	155	def compare(a, b):
	156	if a == b == "": return 0
	157	return self.typ.compare(self.typ.decode(a), self.typ.decode(b))
	158	db.set_flags(bd.DB_DUPSORT)
	159	db.set_bt_compare(compare)
	160	self.bk = db._opendb("i-" + name, bd.DB_BTREE, fl, initdb)
	161	self.bk.set_get_returns_none(False)
	162
	163	def close(self):
	164	self.bk.close()
	165
	166	class cursor(lib.closable):
6efe4e23	167	def __init__(self, idx, fd, fi, ld, li, reverse):
a95055e8	168	self.idx = idx
6efe4e23 FT	169	self.typ = idx.typ
	170	self.cur = self.idx.bk.cursor()
	171	self.item = None
	172	self.fd = fd
	173	self.fi = fi
	174	self.ld = ld
	175	self.li = li
	176	self.rev = reverse
a95055e8 FT	177
	178	def close(self):
	179	if self.cur is not None:
	180	self.cur.close()
6efe4e23	181	self.cur = None
a95055e8 FT	182
	183	def __iter__(self):
	184	return self
	185
6efe4e23 FT	186	def _decode(self, d):
6efe4e23 FT	187	k, v = d
a48a2d5d	188	k = self.typ.decode(k)
6efe4e23 FT	189	v = struct.unpack(">Q", v)[0]
6efe4e23 FT	190	return k, v
a95055e8	191
6efe4e23 FT	192	@dloopfun
	193	def first(self):
	194	try:
d6d41a45	195	if self.fd is missing:
6efe4e23 FT	196	self.item = self._decode(self.cur.first())
	197	else:
	198	k, v = self._decode(self.cur.set_range(self.typ.encode(self.fd)))
	199	if not self.fi:
	200	while self.typ.compare(k, self.fd) == 0:
	201	k, v = self._decode(self.cur.next())
	202	self.item = k, v
	203	except notfound:
	204	self.item = StopIteration
	205
	206	@dloopfun
	207	def last(self):
	208	try:
d6d41a45	209	if self.ld is missing:
6efe4e23 FT	210	self.item = self._decode(self.cur.last())
6efe4e23 FT	211	else:
d6d41a45 FT	212	try:
	213	k, v = self._decode(self.cur.set_range(self.typ.encode(self.ld)))
	214	except notfound:
	215	k, v = self._decode(self.cur.last())
	216	if self.li:
	217	while self.typ.compare(k, self.ld) == 0:
6efe4e23	218	k, v = self._decode(self.cur.next())
d6d41a45 FT	219	while self.typ.compare(k, self.ld) > 0:
d6d41a45 FT	220	k, v = self._decode(self.cur.prev())
6efe4e23	221	else:
d6d41a45	222	while self.typ.compare(k, self.ld) >= 0:
6efe4e23 FT	223	k, v = self._decode(self.cur.prev())
	224	self.item = k, v
	225	except notfound:
	226	self.item = StopIteration
	227
	228	@dloopfun
	229	def next(self):
	230	try:
	231	k, v = self.item = self._decode(self.cur.next())
d6d41a45 FT	232	if (self.ld is not missing and
	233	((self.li and self.typ.compare(k, self.ld) > 0) or
	234	(not self.li and self.typ.compare(k, self.ld) >= 0))):
6efe4e23 FT	235	self.item = StopIteration
	236	except notfound:
	237	self.item = StopIteration
	238
	239	@dloopfun
	240	def prev(self):
a95055e8	241	try:
6efe4e23	242	self.item = self._decode(self.cur.prev())
d6d41a45 FT	243	if (self.fd is not missing and
	244	((self.fi and self.typ.compare(k, self.fd) < 0) or
	245	(not self.fi and self.typ.compare(k, self.fd) <= 0))):
6efe4e23	246	self.item = StopIteration
a95055e8	247	except notfound:
6efe4e23 FT	248	self.item = StopIteration
	249
	250	def __next__(self):
6efe4e23 FT	251	if self.item is None:
	252	if not self.rev:
	253	self.next()
	254	else:
	255	self.prev()
a48a2d5d FT	256	if self.item is StopIteration:
a48a2d5d FT	257	raise StopIteration()
6efe4e23 FT	258	ret, self.item = self.item, None
6efe4e23 FT	259	return ret
a95055e8 FT	260
	261	def skip(self, n=1):
	262	try:
	263	for i in range(n):
	264	next(self)
	265	except StopIteration:
	266	return
	267
6efe4e23 FT	268	def get(self, *, match=missing, ge=missing, gt=missing, lt=missing, le=missing, all=False, reverse=False):
6efe4e23 FT	269	if all:
d6d41a45	270	cur = self.cursor(self, missing, True, missing, True, reverse)
6efe4e23 FT	271	elif match is not missing:
	272	cur = self.cursor(self, match, True, match, True, reverse)
	273	elif ge is not missing or gt is not missing or lt is not missing or le is not missing:
	274	if ge is not missing:
	275	fd, fi = ge, True
	276	elif gt is not missing:
	277	fd, fi = gt, False
	278	else:
d6d41a45	279	fd, fi = missing, True
6efe4e23 FT	280	if le is not missing:
	281	ld, li = le, True
	282	elif lt is not missing:
	283	ld, li = lt, False
	284	else:
d6d41a45	285	ld, li = missing, True
6efe4e23 FT	286	cur = self.cursor(self, fd, fi, ld, li, reverse)
	287	else:
	288	raise NameError("invalid get() specification")
	289	done = False
	290	try:
	291	if not reverse:
	292	cur.first()
	293	else:
	294	cur.last()
	295	done = True
	296	return cur
	297	finally:
	298	if not done:
	299	cur.close()
a95055e8	300
8950191c FT	301	@txnfun(lambda self: self.db.env.env)
	302	def put(self, key, id, *, tx):
	303	obid = struct.pack(">Q", id)
	304	if not self.db.ob.has_key(obid, txn=tx.tx):
	305	raise ValueError("no such object in database: " + str(id))
	306	try:
	307	self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA)
	308	except bd.DBKeyExistError:
	309	return False
	310	return True
	311
	312	@txnfun(lambda self: self.db.env.env)
	313	def remove(self, key, id, *, tx):
	314	obid = struct.pack(">Q", id)
	315	if not self.db.ob.has_key(obid, txn=tx.tx):
	316	raise ValueError("no such object in database: " + str(id))
	317	cur = self.bk.cursor(txn=tx.tx)
	318	try:
a95055e8	319	try:
8950191c FT	320	cur.get_both(self.typ.encode(key), obid)
	321	except notfound:
	322	return False
	323	cur.delete()
	324	finally:
	325	cur.close()
	326	return True