[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_dbid", "t_float", "t_str", "t_casestr", "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 foldtype(simpletype):
    def __init__(self, encode, decode, fold):
        super().__init__(encode, decode)
        self.fold = fold

    def compare(self, a, b):
        return super().compare(self.fold(a), self.fold(b))

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_dbid = t_uint
t_float = simpletype.struct(">d")
t_float.compare = floatcmp
t_str = simpletype((lambda ob: ob.encode("utf-8")), (lambda dat: dat.decode("utf-8")))
t_casestr = foldtype((lambda ob: ob.encode("utf-8")), (lambda dat: dat.decode("utf-8")),
                     (lambda st: st.lower()))

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