11fe17a031fe665c54997fbc765df6082e582643
[didex.git] / 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