1"""
2Vendoring of pickleshare, reduced to used functionalities.
3
4---
5
6PickleShare - a small 'shelve' like datastore with concurrency support
7
8Like shelve, a PickleShareDB object acts like a normal dictionary. Unlike
9shelve, many processes can access the database simultaneously. Changing a
10value in database is immediately visible to other processes accessing the
11same database.
12
13Concurrency is possible because the values are stored in separate files. Hence
14the "database" is a directory where *all* files are governed by PickleShare.
15
16Example usage::
17
18 from pickleshare import *
19 db = PickleShareDB('~/testpickleshare')
20 db.clear()
21 print "Should be empty:",db.items()
22 db['hello'] = 15
23 db['aku ankka'] = [1,2,313]
24 db['paths/are/ok/key'] = [1,(5,46)]
25 print db.keys()
26 del db['aku ankka']
27
28This module is certainly not ZODB, but can be used for low-load
29(non-mission-critical) situations where tiny code size trumps the
30advanced features of a "real" object database.
31
32Installation guide: pip install pickleshare
33
34Author: Ville Vainio <vivainio@gmail.com>
35License: MIT open source license.
36
37"""
38
39__version__ = "0.7.5"
40
41from pathlib import Path
42
43
44import os, stat, time
45
46import collections.abc as collections_abc
47import pickle
48import errno
49
50
51def gethashfile(key):
52 return ("%02x" % abs(hash(key) % 256))[-2:]
53
54
55_sentinel = object()
56
57
58class PickleShareDB(collections_abc.MutableMapping):
59 """The main 'connection' object for PickleShare database"""
60
61 def __init__(self, root):
62 """Return a db object that will manage the specied directory"""
63 if not isinstance(root, str):
64 root = str(root)
65 root = os.path.abspath(os.path.expanduser(root))
66 self.root = Path(root)
67 if not self.root.is_dir():
68 # catching the exception is necessary if multiple processes are concurrently trying to create a folder
69 # exists_ok keyword argument of mkdir does the same but only from Python 3.5
70 try:
71 self.root.mkdir(parents=True)
72 except OSError as e:
73 if e.errno != errno.EEXIST:
74 raise
75 # cache has { 'key' : (obj, orig_mod_time) }
76 self.cache = {}
77
78 def __getitem__(self, key):
79 """db['key'] reading"""
80 fil = self.root / key
81 try:
82 mtime = fil.stat()[stat.ST_MTIME]
83 except OSError:
84 raise KeyError(key)
85
86 if fil in self.cache and mtime == self.cache[fil][1]:
87 return self.cache[fil][0]
88 try:
89 # The cached item has expired, need to read
90 with fil.open("rb") as f:
91 obj = pickle.loads(f.read())
92 except Exception:
93 raise KeyError(key)
94
95 self.cache[fil] = (obj, mtime)
96 return obj
97
98 def __setitem__(self, key, value):
99 """db['key'] = 5"""
100 fil = self.root / key
101 parent = fil.parent
102 if parent and not parent.is_dir():
103 parent.mkdir(parents=True)
104 # We specify protocol 2, so that we can mostly go between Python 2
105 # and Python 3. We can upgrade to protocol 3 when Python 2 is obsolete.
106 with fil.open("wb") as f:
107 pickle.dump(value, f, protocol=2)
108 try:
109 self.cache[fil] = (value, fil.stat().st_mtime)
110 except OSError as e:
111 if e.errno != errno.ENOENT:
112 raise
113
114 def hset(self, hashroot, key, value):
115 """hashed set"""
116 hroot = self.root / hashroot
117 if not hroot.is_dir():
118 hroot.mkdir()
119 hfile = hroot / gethashfile(key)
120 d = self.get(hfile, {})
121 d.update({key: value})
122 self[hfile] = d
123
124 def hget(self, hashroot, key, default=_sentinel, fast_only=True):
125 """hashed get"""
126 hroot = self.root / hashroot
127 hfile = hroot / gethashfile(key)
128
129 d = self.get(hfile, _sentinel)
130 # print "got dict",d,"from",hfile
131 if d is _sentinel:
132 if fast_only:
133 if default is _sentinel:
134 raise KeyError(key)
135
136 return default
137
138 # slow mode ok, works even after hcompress()
139 d = self.hdict(hashroot)
140
141 return d.get(key, default)
142
143 def hdict(self, hashroot):
144 """Get all data contained in hashed category 'hashroot' as dict"""
145 hfiles = self.keys(hashroot + "/*")
146 hfiles.sort()
147 last = len(hfiles) and hfiles[-1] or ""
148 if last.endswith("xx"):
149 # print "using xx"
150 hfiles = [last] + hfiles[:-1]
151
152 all = {}
153
154 for f in hfiles:
155 # print "using",f
156 try:
157 all.update(self[f])
158 except KeyError:
159 print("Corrupt", f, "deleted - hset is not threadsafe!")
160 del self[f]
161
162 self.uncache(f)
163
164 return all
165
166 def hcompress(self, hashroot):
167 """Compress category 'hashroot', so hset is fast again
168
169 hget will fail if fast_only is True for compressed items (that were
170 hset before hcompress).
171
172 """
173 hfiles = self.keys(hashroot + "/*")
174 all = {}
175 for f in hfiles:
176 # print "using",f
177 all.update(self[f])
178 self.uncache(f)
179
180 self[hashroot + "/xx"] = all
181 for f in hfiles:
182 p = self.root / f
183 if p.name == "xx":
184 continue
185 p.unlink()
186
187 def __delitem__(self, key):
188 """del db["key"]"""
189 fil = self.root / key
190 self.cache.pop(fil, None)
191 try:
192 fil.unlink()
193 except OSError:
194 # notfound and permission denied are ok - we
195 # lost, the other process wins the conflict
196 pass
197
198 def _normalized(self, p):
199 """Make a key suitable for user's eyes"""
200 return str(p.relative_to(self.root)).replace("\\", "/")
201
202 def keys(self, globpat=None):
203 """All keys in DB, or all keys matching a glob"""
204
205 if globpat is None:
206 files = self.root.rglob("*")
207 else:
208 files = self.root.glob(globpat)
209 return [self._normalized(p) for p in files if p.is_file()]
210
211 def __iter__(self):
212 return iter(self.keys())
213
214 def __len__(self):
215 return len(self.keys())
216
217 def uncache(self, *items):
218 """Removes all, or specified items from cache
219
220 Use this after reading a large amount of large objects
221 to free up memory, when you won't be needing the objects
222 for a while.
223
224 """
225 if not items:
226 self.cache = {}
227 for it in items:
228 self.cache.pop(it, None)
229
230 def waitget(self, key, maxwaittime=60):
231 """Wait (poll) for a key to get a value
232
233 Will wait for `maxwaittime` seconds before raising a KeyError.
234 The call exits normally if the `key` field in db gets a value
235 within the timeout period.
236
237 Use this for synchronizing different processes or for ensuring
238 that an unfortunately timed "db['key'] = newvalue" operation
239 in another process (which causes all 'get' operation to cause a
240 KeyError for the duration of pickling) won't screw up your program
241 logic.
242 """
243
244 wtimes = [0.2] * 3 + [0.5] * 2 + [1]
245 tries = 0
246 waited = 0
247 while 1:
248 try:
249 val = self[key]
250 return val
251 except KeyError:
252 pass
253
254 if waited > maxwaittime:
255 raise KeyError(key)
256
257 time.sleep(wtimes[tries])
258 waited += wtimes[tries]
259 if tries < len(wtimes) - 1:
260 tries += 1
261
262 def getlink(self, folder):
263 """Get a convenient link for accessing items"""
264 return PickleShareLink(self, folder)
265
266 def __repr__(self):
267 return "PickleShareDB('%s')" % self.root
268
269
270class PickleShareLink:
271 """A shortdand for accessing nested PickleShare data conveniently.
272
273 Created through PickleShareDB.getlink(), example::
274
275 lnk = db.getlink('myobjects/test')
276 lnk.foo = 2
277 lnk.bar = lnk.foo + 5
278
279 """
280
281 def __init__(self, db, keydir):
282 self.__dict__.update(locals())
283
284 def __getattr__(self, key):
285 return self.__dict__["db"][self.__dict__["keydir"] + "/" + key]
286
287 def __setattr__(self, key, val):
288 self.db[self.keydir + "/" + key] = val
289
290 def __repr__(self):
291 db = self.__dict__["db"]
292 keys = db.keys(self.__dict__["keydir"] + "/*")
293 return "<PickleShareLink '{}': {}>".format(
294 self.__dict__["keydir"],
295 ";".join([Path(k).basename() for k in keys]),
296 )