shelve.py 8.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243
  1. """Manage shelves of pickled objects.
  2. A "shelf" is a persistent, dictionary-like object. The difference
  3. with dbm databases is that the values (not the keys!) in a shelf can
  4. be essentially arbitrary Python objects -- anything that the "pickle"
  5. module can handle. This includes most class instances, recursive data
  6. types, and objects containing lots of shared sub-objects. The keys
  7. are ordinary strings.
  8. To summarize the interface (key is a string, data is an arbitrary
  9. object):
  10. import shelve
  11. d = shelve.open(filename) # open, with (g)dbm filename -- no suffix
  12. d[key] = data # store data at key (overwrites old data if
  13. # using an existing key)
  14. data = d[key] # retrieve a COPY of the data at key (raise
  15. # KeyError if no such key) -- NOTE that this
  16. # access returns a *copy* of the entry!
  17. del d[key] # delete data stored at key (raises KeyError
  18. # if no such key)
  19. flag = key in d # true if the key exists
  20. list = d.keys() # a list of all existing keys (slow!)
  21. d.close() # close it
  22. Dependent on the implementation, closing a persistent dictionary may
  23. or may not be necessary to flush changes to disk.
  24. Normally, d[key] returns a COPY of the entry. This needs care when
  25. mutable entries are mutated: for example, if d[key] is a list,
  26. d[key].append(anitem)
  27. does NOT modify the entry d[key] itself, as stored in the persistent
  28. mapping -- it only modifies the copy, which is then immediately
  29. discarded, so that the append has NO effect whatsoever. To append an
  30. item to d[key] in a way that will affect the persistent mapping, use:
  31. data = d[key]
  32. data.append(anitem)
  33. d[key] = data
  34. To avoid the problem with mutable entries, you may pass the keyword
  35. argument writeback=True in the call to shelve.open. When you use:
  36. d = shelve.open(filename, writeback=True)
  37. then d keeps a cache of all entries you access, and writes them all back
  38. to the persistent mapping when you call d.close(). This ensures that
  39. such usage as d[key].append(anitem) works as intended.
  40. However, using keyword argument writeback=True may consume vast amount
  41. of memory for the cache, and it may make d.close() very slow, if you
  42. access many of d's entries after opening it in this way: d has no way to
  43. check which of the entries you access are mutable and/or which ones you
  44. actually mutate, so it must cache, and write back at close, all of the
  45. entries that you access. You can call d.sync() to write back all the
  46. entries in the cache, and empty the cache (d.sync() also synchronizes
  47. the persistent dictionary on disk, if feasible).
  48. """
  49. from pickle import DEFAULT_PROTOCOL, Pickler, Unpickler
  50. from io import BytesIO
  51. import collections.abc
  52. __all__ = ["Shelf", "BsdDbShelf", "DbfilenameShelf", "open"]
  53. class _ClosedDict(collections.abc.MutableMapping):
  54. 'Marker for a closed dict. Access attempts raise a ValueError.'
  55. def closed(self, *args):
  56. raise ValueError('invalid operation on closed shelf')
  57. __iter__ = __len__ = __getitem__ = __setitem__ = __delitem__ = keys = closed
  58. def __repr__(self):
  59. return '<Closed Dictionary>'
  60. class Shelf(collections.abc.MutableMapping):
  61. """Base class for shelf implementations.
  62. This is initialized with a dictionary-like object.
  63. See the module's __doc__ string for an overview of the interface.
  64. """
  65. def __init__(self, dict, protocol=None, writeback=False,
  66. keyencoding="utf-8"):
  67. self.dict = dict
  68. if protocol is None:
  69. protocol = DEFAULT_PROTOCOL
  70. self._protocol = protocol
  71. self.writeback = writeback
  72. self.cache = {}
  73. self.keyencoding = keyencoding
  74. def __iter__(self):
  75. for k in self.dict.keys():
  76. yield k.decode(self.keyencoding)
  77. def __len__(self):
  78. return len(self.dict)
  79. def __contains__(self, key):
  80. return key.encode(self.keyencoding) in self.dict
  81. def get(self, key, default=None):
  82. if key.encode(self.keyencoding) in self.dict:
  83. return self[key]
  84. return default
  85. def __getitem__(self, key):
  86. try:
  87. value = self.cache[key]
  88. except KeyError:
  89. f = BytesIO(self.dict[key.encode(self.keyencoding)])
  90. value = Unpickler(f).load()
  91. if self.writeback:
  92. self.cache[key] = value
  93. return value
  94. def __setitem__(self, key, value):
  95. if self.writeback:
  96. self.cache[key] = value
  97. f = BytesIO()
  98. p = Pickler(f, self._protocol)
  99. p.dump(value)
  100. self.dict[key.encode(self.keyencoding)] = f.getvalue()
  101. def __delitem__(self, key):
  102. del self.dict[key.encode(self.keyencoding)]
  103. try:
  104. del self.cache[key]
  105. except KeyError:
  106. pass
  107. def __enter__(self):
  108. return self
  109. def __exit__(self, type, value, traceback):
  110. self.close()
  111. def close(self):
  112. if self.dict is None:
  113. return
  114. try:
  115. self.sync()
  116. try:
  117. self.dict.close()
  118. except AttributeError:
  119. pass
  120. finally:
  121. # Catch errors that may happen when close is called from __del__
  122. # because CPython is in interpreter shutdown.
  123. try:
  124. self.dict = _ClosedDict()
  125. except:
  126. self.dict = None
  127. def __del__(self):
  128. if not hasattr(self, 'writeback'):
  129. # __init__ didn't succeed, so don't bother closing
  130. # see http://bugs.python.org/issue1339007 for details
  131. return
  132. self.close()
  133. def sync(self):
  134. if self.writeback and self.cache:
  135. self.writeback = False
  136. for key, entry in self.cache.items():
  137. self[key] = entry
  138. self.writeback = True
  139. self.cache = {}
  140. if hasattr(self.dict, 'sync'):
  141. self.dict.sync()
  142. class BsdDbShelf(Shelf):
  143. """Shelf implementation using the "BSD" db interface.
  144. This adds methods first(), next(), previous(), last() and
  145. set_location() that have no counterpart in [g]dbm databases.
  146. The actual database must be opened using one of the "bsddb"
  147. modules "open" routines (i.e. bsddb.hashopen, bsddb.btopen or
  148. bsddb.rnopen) and passed to the constructor.
  149. See the module's __doc__ string for an overview of the interface.
  150. """
  151. def __init__(self, dict, protocol=None, writeback=False,
  152. keyencoding="utf-8"):
  153. Shelf.__init__(self, dict, protocol, writeback, keyencoding)
  154. def set_location(self, key):
  155. (key, value) = self.dict.set_location(key)
  156. f = BytesIO(value)
  157. return (key.decode(self.keyencoding), Unpickler(f).load())
  158. def next(self):
  159. (key, value) = next(self.dict)
  160. f = BytesIO(value)
  161. return (key.decode(self.keyencoding), Unpickler(f).load())
  162. def previous(self):
  163. (key, value) = self.dict.previous()
  164. f = BytesIO(value)
  165. return (key.decode(self.keyencoding), Unpickler(f).load())
  166. def first(self):
  167. (key, value) = self.dict.first()
  168. f = BytesIO(value)
  169. return (key.decode(self.keyencoding), Unpickler(f).load())
  170. def last(self):
  171. (key, value) = self.dict.last()
  172. f = BytesIO(value)
  173. return (key.decode(self.keyencoding), Unpickler(f).load())
  174. class DbfilenameShelf(Shelf):
  175. """Shelf implementation using the "dbm" generic dbm interface.
  176. This is initialized with the filename for the dbm database.
  177. See the module's __doc__ string for an overview of the interface.
  178. """
  179. def __init__(self, filename, flag='c', protocol=None, writeback=False):
  180. import dbm
  181. Shelf.__init__(self, dbm.open(filename, flag), protocol, writeback)
  182. def open(filename, flag='c', protocol=None, writeback=False):
  183. """Open a persistent dictionary for reading and writing.
  184. The filename parameter is the base filename for the underlying
  185. database. As a side-effect, an extension may be added to the
  186. filename and more than one file may be created. The optional flag
  187. parameter has the same interpretation as the flag parameter of
  188. dbm.open(). The optional protocol parameter specifies the
  189. version of the pickle protocol.
  190. See the module's __doc__ string for an overview of the interface.
  191. """
  192. return DbfilenameShelf(filename, flag, protocol, writeback)