didex/store.py

   1 import threading, pickle
   2 from . import db, index, cache
   3 from .db import txnfun
   4
   5 class environment(object):
   6     def __init__(self, path):
   7         self.path = path
   8         self.lk = threading.Lock()
   9         self.bk = None
  10
  11     def __call__(self):
  12         with self.lk:
  13             if self.bk is None:
  14                 self.bk = db.environment(self.path)
  15             return self.bk
  16
  17     def close(self):
  18         with self.lk:
  19             if self.bk is not None:
  20                 self.bk.close()
  21                 self.bk = None
  22
  23 class storedesc(object):
  24     pass
  25
  26 def storedescs(obj):
  27     t = type(obj)
  28     ret = getattr(t, "__didex_attr", None)
  29     if ret is None:
  30         ret = []
  31         for nm, val in t.__dict__.items():
  32             if isinstance(val, storedesc):
  33                 ret.append((nm, val))
  34         t.__didex_attr = ret
  35     return ret
  36
  37 class store(object):
  38     def __init__(self, name, *, env=None, path=".", ncache=None):
  39         self.name = name
  40         self.lk = threading.Lock()
  41         if env:
  42             self.env = env
  43         else:
  44             self.env = environment(path)
  45         self._db = None
  46         if ncache is None:
  47             ncache = cache.cache()
  48         self.cache = ncache
  49         self.cache.load = self._load
  50
  51     def db(self):
  52         with self.lk:
  53             if self._db is None:
  54                 self._db = self.env().db(self.name)
  55             return self._db
  56
  57     def _load(self, id):
  58         try:
  59             return pickle.loads(self.db().get(id))
  60         except:
  61             raise KeyError(id, "could not unpickle data")
  62
  63     def _encode(self, obj):
  64         return pickle.dumps(obj)
  65
  66     def get(self, id, *, load=True):
  67         return self.cache.get(id, load=load)
  68
  69     @txnfun(lambda self: self.db().env.env)
  70     def register(self, obj, *, tx):
  71         id = self.db().add(self._encode(obj), tx=tx)
  72         for nm, attr in storedescs(obj):
  73             attr.register(id, obj, tx)
  74         self.cache.put(id, obj)
  75         return id
  76
  77     @txnfun(lambda self: self.db().env.env)
  78     def unregister(self, id, *, tx):
  79         obj = self.get(id)
  80         for nm, attr in storedescs(obj):
  81             attr.unregister(id, obj, tx)
  82         self.db().remove(id, tx=tx)
  83         self.cache.remove(id)
  84
  85     @txnfun(lambda self: self.db().env.env)
  86     def update(self, id, *, tx):
  87         obj = self.get(id, load=False)
  88         for nm, attr, in storedescs(obj):
  89             attr.update(id, obj, tx)
  90         self.db().replace(id, self._encode(obj), tx=tx)