Typo fix.
[didex.git] / didex / store.py
CommitLineData
e38ebdef 1import threading, pickle, inspect, atexit, weakref
b080a59c
FT
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)
a19ad473 25 atexit.register(self.close)
b080a59c
FT
26 return self.bk
27
28 def close(self):
29 with self.lk:
30 if self.bk is not None:
a19ad473 31 atexit.unregister(self.close)
b080a59c
FT
32 self.bk.close()
33 self.bk = None
34
35class storedesc(object):
36 pass
37
38def storedescs(obj):
39 t = type(obj)
36c0a011 40 ret = t.__dict__.get("__didex_attr")
b080a59c
FT
41 if ret is None:
42 ret = []
36c0a011
FT
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))
de73859d 47 t.__didex_attr = ret
b080a59c
FT
48 return ret
49
e38ebdef
FT
50class 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
6e81ddd5 72class datastore(object):
e38ebdef 73 def __init__(self, name, *, env=None, path=".", ncache=None, codec=None):
b080a59c
FT
74 self.name = name
75 self.lk = threading.Lock()
76 if env:
77 self.env = env
78 else:
eca9b3be 79 self.env = environment(path=path)
b080a59c
FT
80 self._db = None
81 if ncache is None:
82 ncache = cache.cache()
e38ebdef
FT
83 if codec is not None:
84 self._encode, self._decode = codec
b080a59c
FT
85 self.cache = ncache
86 self.cache.load = self._load
e38ebdef 87 self.icache = icache()
b080a59c
FT
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
e38ebdef 95 def _decode(self, data):
b080a59c 96 try:
e38ebdef 97 return pickle.loads(data)
b080a59c
FT
98 except:
99 raise KeyError(id, "could not unpickle data")
100
101 def _encode(self, obj):
102 return pickle.dumps(obj)
103
e38ebdef
FT
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
b080a59c
FT
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)
ca180faa 125 def unregister(self, id, *, vfy=None, tx):
b080a59c 126 obj = self.get(id)
ca180faa
FT
127 if vfy is not None and obj is not vfy:
128 raise RuntimeError("object identity crisis: " + str(vfy) + " is not cached object " + obj)
b080a59c
FT
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)
ca180faa 135 def update(self, id, *, vfy=None, tx):
b080a59c 136 obj = self.get(id, load=False)
ca180faa
FT
137 if vfy is not None and obj is not vfy:
138 raise RuntimeError("object identity crisis: " + str(vfy) + " is not cached object " + obj)
b080a59c
FT
139 for nm, attr, in storedescs(obj):
140 attr.update(id, obj, tx)
141 self.db().replace(id, self._encode(obj), tx=tx)
ca180faa
FT
142
143class autotype(type):
144 def __call__(self, *args, **kwargs):
145 new = super().__call__(*args, **kwargs)
146 new.id = self.store.register(new)
e38ebdef 147 # XXX? ID is not saved now, but relied upon to be __didex_loaded__ later.
ca180faa
FT
148 return new
149
150class autostore(object, metaclass=autotype):
151 def __init__(self):
152 self.id = None
153
e38ebdef
FT
154 def __didex_loaded__(self, store, id):
155 assert self.id is None or self.id == id
156 self.id = id
157
ca180faa
FT
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