Commit | Line | Data |
---|---|---|
e38ebdef | 1 | import threading, pickle, inspect, atexit, weakref |
b080a59c FT |
2 | from . import db, index, cache |
3 | from .db import txnfun | |
4 | ||
6e81ddd5 | 5 | __all__ = ["environment", "datastore", "autostore"] |
cbf73d3a | 6 | |
b080a59c | 7 | class 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 | ||
35 | class storedesc(object): | |
36 | pass | |
37 | ||
38 | def 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 |
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 | ||
6e81ddd5 | 72 | class 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 | |
143 | class 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 | ||
150 | class 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 |