didex/store.py

   1 import threading, pickle, inspect
   2 from . import db, index, cache
   3 from .db import txnfun
   4
   5 __all__ = ["environment", "datastore", "autostore"]
   6
   7 class environment(object):
   8     def __init__(self, *, path=None, getpath=None, recover=False):
   9         if path is not None:
  10             self.path = path
  11             self.getpath = None
  12         else:
  13             self.path = None
  14             self.getpath = getpath
  15         self.recover = recover
  16         self.lk = threading.Lock()
  17         self.bk = None
  18
  19     def __call__(self):
  20         with self.lk:
  21             if self.bk is None:
  22                 if self.path is None:
  23                     self.path = self.getpath()
  24                 self.bk = db.environment(self.path, recover=self.recover)
  25             return self.bk
  26
  27     def close(self):
  28         with self.lk:
  29             if self.bk is not None:
  30                 self.bk.close()
  31                 self.bk = None
  32
  33 class storedesc(object):
  34     pass
  35
  36 def storedescs(obj):
  37     t = type(obj)
  38     ret = t.__dict__.get("__didex_attr")
  39     if ret is None:
  40         ret = []
  41         for st in inspect.getmro(t):
  42             for nm, val in st.__dict__.items():
  43                 if isinstance(val, storedesc):
  44                     ret.append((nm, val))
  45         t.__dict__["__didex_attr"] = ret
  46     return ret
  47
  48 class datastore(object):
  49     def __init__(self, name, *, env=None, path=".", ncache=None):
  50         self.name = name
  51         self.lk = threading.Lock()
  52         if env:
  53             self.env = env
  54         else:
  55             self.env = environment(path=path)
  56         self._db = None
  57         if ncache is None:
  58             ncache = cache.cache()
  59         self.cache = ncache
  60         self.cache.load = self._load
  61
  62     def db(self):
  63         with self.lk:
  64             if self._db is None:
  65                 self._db = self.env().db(self.name)
  66             return self._db
  67
  68     def _load(self, id):
  69         try:
  70             return pickle.loads(self.db().get(id))
  71         except:
  72             raise KeyError(id, "could not unpickle data")
  73
  74     def _encode(self, obj):
  75         return pickle.dumps(obj)
  76
  77     def get(self, id, *, load=True):
  78         return self.cache.get(id, load=load)
  79
  80     @txnfun(lambda self: self.db().env.env)
  81     def register(self, obj, *, tx):
  82         id = self.db().add(self._encode(obj), tx=tx)
  83         for nm, attr in storedescs(obj):
  84             attr.register(id, obj, tx)
  85         self.cache.put(id, obj)
  86         return id
  87
  88     @txnfun(lambda self: self.db().env.env)
  89     def unregister(self, id, *, vfy=None, tx):
  90         obj = self.get(id)
  91         if vfy is not None and obj is not vfy:
  92             raise RuntimeError("object identity crisis: " + str(vfy) + " is not cached object " + obj)
  93         for nm, attr in storedescs(obj):
  94             attr.unregister(id, obj, tx)
  95         self.db().remove(id, tx=tx)
  96         self.cache.remove(id)
  97
  98     @txnfun(lambda self: self.db().env.env)
  99     def update(self, id, *, vfy=None, tx):
 100         obj = self.get(id, load=False)
 101         if vfy is not None and obj is not vfy:
 102             raise RuntimeError("object identity crisis: " + str(vfy) + " is not cached object " + obj)
 103         for nm, attr, in storedescs(obj):
 104             attr.update(id, obj, tx)
 105         self.db().replace(id, self._encode(obj), tx=tx)
 106
 107 class autotype(type):
 108     def __call__(self, *args, **kwargs):
 109         new = super().__call__(*args, **kwargs)
 110         new.id = self.store.register(new)
 111         self.store.update(new.id, vfy=new) # This doesn't feel too nice.
 112         return new
 113
 114 class autostore(object, metaclass=autotype):
 115     def __init__(self):
 116         self.id = None
 117
 118     def save(self):
 119         self.store.update(self.id, vfy=self)
 120
 121     def remove(self):
 122         self.store.unregister(self.id, vfy=self)
 123         self.id = None