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