Added convenience get() wrappers.

[didex.git] / didex / values.py

import threading
from . import store, lib, index
from .store import storedesc

__all__ = ["simple", "multi", "compound"]

class cursor(lib.closable):
    def __init__(self, bk, st):
        self.bk = bk
        self.st = st

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

    def __iter__(self):
        return self

    def __next__(self):
        k, id = next(self.bk)
        return k, self.st.get(id)

    def skip(self, n=1):
        self.bk.skip(n)

class base(storedesc):
    def __init__(self, store, indextype, name, datatype):
        self.store = store
        self.indextype = indextype
        self.name = name
        self.typ = datatype
        self.idx = None
        self.lk = threading.Lock()

    def index(self):
        with self.lk:
            if self.idx is None:
                self.idx = self.indextype(self.store.db(), self.name, self.typ)
            return self.idx

    def get(self, **kwargs):
        return cursor(self.index().get(**kwargs), self.store)

    def get1(self, *, check=False, default=KeyError, **kwargs):
        with self.get(**kwargs) as cursor:
            try:
                k, v = next(cursor)
            except StopIteration:
                if default is not KeyError:
                    return default
                raise KeyError("no matches in " + self.name, kwargs)
            if check:
                try:
                    next(cursor)
                except StopIteration:
                    pass
                else:
                    raise ValueError("unexpected multiple matchies in " + self.name, kwargs)
            return v

    def list(self, **kwargs):
        with self.get(**kwargs) as cursor:
            return [v for k, v in cursor]

class descbase(base):
    def __init__(self, store, indextype, name, datatype, default):
        super().__init__(store, indextype, name, datatype)
        self.default = default
        self.mattr = "__idx_%s_new" % name
        self.iattr = "__idx_%s_cur" % name

    def __get__(self, obj, cls):
        if obj is None: return self
        return getattr(obj, self.mattr, self.default)

    def __set__(self, obj, val):
        setattr(obj, self.mattr, val)

    def __delete__(self, obj):
        delattr(obj, self.mattr)

class simple(descbase):
    def __init__(self, store, indextype, name, datatype, default=None):
        super().__init__(store, indextype, name, datatype, default)

    def register(self, id, obj, tx):
        val = self.__get__(obj, None)
        self.index().put(val, id, tx=tx)
        tx.postcommit(lambda: setattr(obj, self.iattr, val))

    def unregister(self, id, obj, tx):
        self.index().remove(getattr(obj, self.iattr), id, tx=tx)
        tx.postcommit(lambda: delattr(obj, self.iattr))

    def update(self, id, obj, tx):
        val = self.__get__(obj, None)
        ival = getattr(obj, self.iattr)
        if val != ival:
            idx = self.index()
            idx.remove(ival, id, tx=tx)
            idx.put(val, id, tx=tx)
            tx.postcommit(lambda: setattr(obj, self.iattr, val))

class multi(descbase):
    def __init__(self, store, indextype, name, datatype):
        super().__init__(store, indextype, name, datatype, ())

    def register(self, id, obj, tx):
        vals = frozenset(self.__get__(obj, None))
        idx = self.index()
        for val in vals:
            idx.put(val, id, tx=tx)
        tx.postcommit(lambda: setattr(obj, self.iattr, vals))

    def unregister(self, id, obj, tx):
        idx = self.index()
        for val in getattr(obj, self.iattr):
            idx.remove(val, id, tx=tx)
        tx.postcommit(lambda: delattr(obj, self.iattr))

    def update(self, id, obj, tx):
        vals = frozenset(self.__get__(obj, None))
        ivals = getattr(obj, self.iattr)
        if vals != ivals:
            idx = self.index()
            for val in ivals - vals:
                idx.remove(val, id, tx=tx)
            for val in vals - ivals:
                idx.put(val, id, tx=tx)
            tx.postcommit(lambda: setattr(obj, self.iattr, vals))

class compound(base):
    def __init__(self, indextype, name, *parts):
        super().__init__(parts[0].store, indextype, name, index.compound(*(part.typ for part in parts)))
        self.parts = parts
        self.iattr = "__idx_%s_cur" % name

    def minim(self, *parts):
        return self.typ.minim(*parts)
    def maxim(self, *parts):
        return self.typ.maxim(*parts)

    def get(self, *, partial=None, **spec):
        if partial is not None:
            return super().get(ge=self.minim(*partial), le = self.maxim(*partial), **spec)
        else:
            return super().get(**spec)

    def register(self, id, obj, tx):
        val = tuple(part.__get__(obj, None) for part in self.parts)
        self.index().put(val, id, tx=tx)
        tx.postcommit(lambda: setattr(obj, self.iattr, val))

    def unregister(self, id, obj, tx):
        self.index().remove(getattr(obj, self.iattr), id, tx=tx)
        tx.postcommit(lambda: delattr(obj, self.iattr))

    def update(self, id, obj, tx):
        val = tuple(part.__get__(obj, None) for part in self.parts)
        ival = getattr(obj, self.iattr)
        if val != ival:
            idx = self.index()
            idx.remove(ival, id, tx=tx)
            idx.put(val, id, tx=tx)
            tx.postcommit(lambda: setattr(obj, self.iattr, val))