didex/store.py

   1 import threading, pickle, inspect, atexit, weakref
   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                 atexit.register(self.close)
  26             return self.bk
  27
  28     def close(self):
  29         with self.lk:
  30             if self.bk is not None:
  31                 atexit.unregister(self.close)
  32                 self.bk.close()
  33                 self.bk = None
  34
  35 class storedesc(object):
  36     pass
  37
  38 def storedescs(obj):
  39     t = type(obj)
  40     ret = t.__dict__.get("__didex_attr")
  41     if ret is None:
  42         ret = []
  43         for st in inspect.getmro(t):
  44             for nm, val in st.__dict__.items():
  45                 if isinstance(val, storedesc):
  46                     ret.append((nm, val))
  47         t.__didex_attr = ret
  48     return ret
  49
  50 class icache(object):
  51     def __init__(self):
  52         self.d = weakref.WeakKeyDictionary()
  53
  54     def __getitem__(self, key):
  55         obj, idx = key
  56         return self.d[obj][idx]
  57     def __setitem__(self, key, val):
  58         obj, idx = key
  59         if obj in self.d:
  60             self.d[obj][idx] = val
  61         else:
  62             self.d[obj] = {idx: val}
  63     def __delitem__(self, key):
  64         obj, idx = key
  65         del self.d[obj][idx]
  66     def get(self, key, default=None):
  67         obj, idx = key
  68         if obj not in self.d:
  69             return default
  70         return self.d[obj].get(idx, default)
  71
  72 class datastore(object):
  73     def __init__(self, name, *, env=None, path=".", ncache=None, codec=None):
  74         self.name = name
  75         self.lk = threading.Lock()
  76         if env:
  77             self.env = env
  78         else:
  79             self.env = environment(path=path)
  80         self._db = None
  81         if ncache is None:
  82             ncache = cache.cache()
  83         if codec is not None:
  84             self._encode, self._decode = codec
  85         self.cache = ncache
  86         self.cache.load = self._load
  87         self.icache = icache()
  88
  89     def db(self):
  90         with self.lk:
  91             if self._db is None:
  92                 self._db = self.env().db(self.name)
  93             return self._db
  94
  95     def _decode(self, data):
  96         try:
  97             return pickle.loads(data)
  98         except:
  99             raise KeyError(id, "could not unpickle data")
 100
 101     def _encode(self, obj):
 102         return pickle.dumps(obj)
 103
 104     @txnfun(lambda self: self.db().env.env)
 105     def _load(self, id, *, tx):
 106         loaded = self._decode(self.db().get(id, tx=tx))
 107         if hasattr(loaded, "__didex_loaded__"):
 108             loaded.__didex_loaded__(self, id)
 109         for nm, attr in storedescs(loaded):
 110             attr.loaded(id, loaded, tx)
 111         return loaded
 112
 113     def get(self, id, *, load=True):
 114         return self.cache.get(id, load=load)
 115
 116     @txnfun(lambda self: self.db().env.env)
 117     def register(self, obj, *, tx):
 118         id = self.db().add(self._encode(obj), tx=tx)
 119         for nm, attr in storedescs(obj):
 120             attr.register(id, obj, tx)
 121         self.cache.put(id, obj)
 122         return id
 123
 124     @txnfun(lambda self: self.db().env.env)
 125     def unregister(self, id, *, vfy=None, tx):
 126         obj = self.get(id)
 127         if vfy is not None and obj is not vfy:
 128             raise RuntimeError("object identity crisis: " + str(vfy) + " is not cached object " + obj)
 129         for nm, attr in storedescs(obj):
 130             attr.unregister(id, obj, tx)
 131         self.db().remove(id, tx=tx)
 132         self.cache.remove(id)
 133
 134     @txnfun(lambda self: self.db().env.env)
 135     def update(self, id, *, vfy=None, tx):
 136         obj = self.get(id, load=False)
 137         if vfy is not None and obj is not vfy:
 138             raise RuntimeError("object identity crisis: " + str(vfy) + " is not cached object " + obj)
 139         for nm, attr, in storedescs(obj):
 140             attr.update(id, obj, tx)
 141         self.db().replace(id, self._encode(obj), tx=tx)
 142
 143 class autotype(type):
 144     def __call__(self, *args, **kwargs):
 145         new = super().__call__(*args, **kwargs)
 146         new.id = self.store.register(new)
 147         # XXX? ID is not saved now, but relied upon to be __didex_loaded__ later.
 148         return new
 149
 150 class autostore(object, metaclass=autotype):
 151     def __init__(self):
 152         self.id = None
 153
 154     def __didex_loaded__(self, store, id):
 155         assert self.id is None or self.id == id
 156         self.id = id
 157
 158     def save(self):
 159         self.store.update(self.id, vfy=self)
 160
 161     def remove(self):
 162         self.store.unregister(self.id, vfy=self)
 163         self.id = None