dumb.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317
  1. """A dumb and slow but simple dbm clone.
  2. For database spam, spam.dir contains the index (a text file),
  3. spam.bak *may* contain a backup of the index (also a text file),
  4. while spam.dat contains the data (a binary file).
  5. XXX TO DO:
  6. - seems to contain a bug when updating...
  7. - reclaim free space (currently, space once occupied by deleted or expanded
  8. items is never reused)
  9. - support concurrent access (currently, if two processes take turns making
  10. updates, they can mess up the index)
  11. - support efficient access to large databases (currently, the whole index
  12. is read when the database is opened, and some updates rewrite the whole index)
  13. - support opening for read-only (flag = 'm')
  14. """
  15. import ast as _ast
  16. import io as _io
  17. import os as _os
  18. import collections.abc
  19. __all__ = ["error", "open"]
  20. _BLOCKSIZE = 512
  21. error = OSError
  22. class _Database(collections.abc.MutableMapping):
  23. # The on-disk directory and data files can remain in mutually
  24. # inconsistent states for an arbitrarily long time (see comments
  25. # at the end of __setitem__). This is only repaired when _commit()
  26. # gets called. One place _commit() gets called is from __del__(),
  27. # and if that occurs at program shutdown time, module globals may
  28. # already have gotten rebound to None. Since it's crucial that
  29. # _commit() finish successfully, we can't ignore shutdown races
  30. # here, and _commit() must not reference any globals.
  31. _os = _os # for _commit()
  32. _io = _io # for _commit()
  33. def __init__(self, filebasename, mode, flag='c'):
  34. filebasename = self._os.fsencode(filebasename)
  35. self._mode = mode
  36. self._readonly = (flag == 'r')
  37. # The directory file is a text file. Each line looks like
  38. # "%r, (%d, %d)\n" % (key, pos, siz)
  39. # where key is the string key, pos is the offset into the dat
  40. # file of the associated value's first byte, and siz is the number
  41. # of bytes in the associated value.
  42. self._dirfile = filebasename + b'.dir'
  43. # The data file is a binary file pointed into by the directory
  44. # file, and holds the values associated with keys. Each value
  45. # begins at a _BLOCKSIZE-aligned byte offset, and is a raw
  46. # binary 8-bit string value.
  47. self._datfile = filebasename + b'.dat'
  48. self._bakfile = filebasename + b'.bak'
  49. # The index is an in-memory dict, mirroring the directory file.
  50. self._index = None # maps keys to (pos, siz) pairs
  51. # Handle the creation
  52. self._create(flag)
  53. self._update(flag)
  54. def _create(self, flag):
  55. if flag == 'n':
  56. for filename in (self._datfile, self._bakfile, self._dirfile):
  57. try:
  58. _os.remove(filename)
  59. except OSError:
  60. pass
  61. # Mod by Jack: create data file if needed
  62. try:
  63. f = _io.open(self._datfile, 'r', encoding="Latin-1")
  64. except OSError:
  65. if flag not in ('c', 'n'):
  66. raise
  67. with _io.open(self._datfile, 'w', encoding="Latin-1") as f:
  68. self._chmod(self._datfile)
  69. else:
  70. f.close()
  71. # Read directory file into the in-memory index dict.
  72. def _update(self, flag):
  73. self._modified = False
  74. self._index = {}
  75. try:
  76. f = _io.open(self._dirfile, 'r', encoding="Latin-1")
  77. except OSError:
  78. if flag not in ('c', 'n'):
  79. raise
  80. self._modified = True
  81. else:
  82. with f:
  83. for line in f:
  84. line = line.rstrip()
  85. key, pos_and_siz_pair = _ast.literal_eval(line)
  86. key = key.encode('Latin-1')
  87. self._index[key] = pos_and_siz_pair
  88. # Write the index dict to the directory file. The original directory
  89. # file (if any) is renamed with a .bak extension first. If a .bak
  90. # file currently exists, it's deleted.
  91. def _commit(self):
  92. # CAUTION: It's vital that _commit() succeed, and _commit() can
  93. # be called from __del__(). Therefore we must never reference a
  94. # global in this routine.
  95. if self._index is None or not self._modified:
  96. return # nothing to do
  97. try:
  98. self._os.unlink(self._bakfile)
  99. except OSError:
  100. pass
  101. try:
  102. self._os.rename(self._dirfile, self._bakfile)
  103. except OSError:
  104. pass
  105. with self._io.open(self._dirfile, 'w', encoding="Latin-1") as f:
  106. self._chmod(self._dirfile)
  107. for key, pos_and_siz_pair in self._index.items():
  108. # Use Latin-1 since it has no qualms with any value in any
  109. # position; UTF-8, though, does care sometimes.
  110. entry = "%r, %r\n" % (key.decode('Latin-1'), pos_and_siz_pair)
  111. f.write(entry)
  112. sync = _commit
  113. def _verify_open(self):
  114. if self._index is None:
  115. raise error('DBM object has already been closed')
  116. def __getitem__(self, key):
  117. if isinstance(key, str):
  118. key = key.encode('utf-8')
  119. self._verify_open()
  120. pos, siz = self._index[key] # may raise KeyError
  121. with _io.open(self._datfile, 'rb') as f:
  122. f.seek(pos)
  123. dat = f.read(siz)
  124. return dat
  125. # Append val to the data file, starting at a _BLOCKSIZE-aligned
  126. # offset. The data file is first padded with NUL bytes (if needed)
  127. # to get to an aligned offset. Return pair
  128. # (starting offset of val, len(val))
  129. def _addval(self, val):
  130. with _io.open(self._datfile, 'rb+') as f:
  131. f.seek(0, 2)
  132. pos = int(f.tell())
  133. npos = ((pos + _BLOCKSIZE - 1) // _BLOCKSIZE) * _BLOCKSIZE
  134. f.write(b'\0'*(npos-pos))
  135. pos = npos
  136. f.write(val)
  137. return (pos, len(val))
  138. # Write val to the data file, starting at offset pos. The caller
  139. # is responsible for ensuring that there's enough room starting at
  140. # pos to hold val, without overwriting some other value. Return
  141. # pair (pos, len(val)).
  142. def _setval(self, pos, val):
  143. with _io.open(self._datfile, 'rb+') as f:
  144. f.seek(pos)
  145. f.write(val)
  146. return (pos, len(val))
  147. # key is a new key whose associated value starts in the data file
  148. # at offset pos and with length siz. Add an index record to
  149. # the in-memory index dict, and append one to the directory file.
  150. def _addkey(self, key, pos_and_siz_pair):
  151. self._index[key] = pos_and_siz_pair
  152. with _io.open(self._dirfile, 'a', encoding="Latin-1") as f:
  153. self._chmod(self._dirfile)
  154. f.write("%r, %r\n" % (key.decode("Latin-1"), pos_and_siz_pair))
  155. def __setitem__(self, key, val):
  156. if self._readonly:
  157. raise error('The database is opened for reading only')
  158. if isinstance(key, str):
  159. key = key.encode('utf-8')
  160. elif not isinstance(key, (bytes, bytearray)):
  161. raise TypeError("keys must be bytes or strings")
  162. if isinstance(val, str):
  163. val = val.encode('utf-8')
  164. elif not isinstance(val, (bytes, bytearray)):
  165. raise TypeError("values must be bytes or strings")
  166. self._verify_open()
  167. self._modified = True
  168. if key not in self._index:
  169. self._addkey(key, self._addval(val))
  170. else:
  171. # See whether the new value is small enough to fit in the
  172. # (padded) space currently occupied by the old value.
  173. pos, siz = self._index[key]
  174. oldblocks = (siz + _BLOCKSIZE - 1) // _BLOCKSIZE
  175. newblocks = (len(val) + _BLOCKSIZE - 1) // _BLOCKSIZE
  176. if newblocks <= oldblocks:
  177. self._index[key] = self._setval(pos, val)
  178. else:
  179. # The new value doesn't fit in the (padded) space used
  180. # by the old value. The blocks used by the old value are
  181. # forever lost.
  182. self._index[key] = self._addval(val)
  183. # Note that _index may be out of synch with the directory
  184. # file now: _setval() and _addval() don't update the directory
  185. # file. This also means that the on-disk directory and data
  186. # files are in a mutually inconsistent state, and they'll
  187. # remain that way until _commit() is called. Note that this
  188. # is a disaster (for the database) if the program crashes
  189. # (so that _commit() never gets called).
  190. def __delitem__(self, key):
  191. if self._readonly:
  192. raise error('The database is opened for reading only')
  193. if isinstance(key, str):
  194. key = key.encode('utf-8')
  195. self._verify_open()
  196. self._modified = True
  197. # The blocks used by the associated value are lost.
  198. del self._index[key]
  199. # XXX It's unclear why we do a _commit() here (the code always
  200. # XXX has, so I'm not changing it). __setitem__ doesn't try to
  201. # XXX keep the directory file in synch. Why should we? Or
  202. # XXX why shouldn't __setitem__?
  203. self._commit()
  204. def keys(self):
  205. try:
  206. return list(self._index)
  207. except TypeError:
  208. raise error('DBM object has already been closed') from None
  209. def items(self):
  210. self._verify_open()
  211. return [(key, self[key]) for key in self._index.keys()]
  212. def __contains__(self, key):
  213. if isinstance(key, str):
  214. key = key.encode('utf-8')
  215. try:
  216. return key in self._index
  217. except TypeError:
  218. if self._index is None:
  219. raise error('DBM object has already been closed') from None
  220. else:
  221. raise
  222. def iterkeys(self):
  223. try:
  224. return iter(self._index)
  225. except TypeError:
  226. raise error('DBM object has already been closed') from None
  227. __iter__ = iterkeys
  228. def __len__(self):
  229. try:
  230. return len(self._index)
  231. except TypeError:
  232. raise error('DBM object has already been closed') from None
  233. def close(self):
  234. try:
  235. self._commit()
  236. finally:
  237. self._index = self._datfile = self._dirfile = self._bakfile = None
  238. __del__ = close
  239. def _chmod(self, file):
  240. self._os.chmod(file, self._mode)
  241. def __enter__(self):
  242. return self
  243. def __exit__(self, *args):
  244. self.close()
  245. def open(file, flag='c', mode=0o666):
  246. """Open the database file, filename, and return corresponding object.
  247. The flag argument, used to control how the database is opened in the
  248. other DBM implementations, supports only the semantics of 'c' and 'n'
  249. values. Other values will default to the semantics of 'c' value:
  250. the database will always opened for update and will be created if it
  251. does not exist.
  252. The optional mode argument is the UNIX mode of the file, used only when
  253. the database has to be created. It defaults to octal code 0o666 (and
  254. will be modified by the prevailing umask).
  255. """
  256. # Modify mode depending on the umask
  257. try:
  258. um = _os.umask(0)
  259. _os.umask(um)
  260. except AttributeError:
  261. pass
  262. else:
  263. # Turn off any bits that are set in the umask
  264. mode = mode & (~um)
  265. if flag not in ('r', 'w', 'c', 'n'):
  266. raise ValueError("Flag must be one of 'r', 'w', 'c', or 'n'")
  267. return _Database(file, mode, flag=flag)