pickleshare.py 8.6 KB

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