dbshelve.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381
  1. #------------------------------------------------------------------------
  2. # Copyright (c) 1997-2001 by Total Control Software
  3. # All Rights Reserved
  4. #------------------------------------------------------------------------
  5. #
  6. # Module Name: dbShelve.py
  7. #
  8. # Description: A reimplementation of the standard shelve.py that
  9. # forces the use of cPickle, and DB.
  10. #
  11. # Creation Date: 11/3/97 3:39:04PM
  12. #
  13. # License: This is free software. You may use this software for any
  14. # purpose including modification/redistribution, so long as
  15. # this header remains intact and that you do not claim any
  16. # rights of ownership or authorship of this software. This
  17. # software has been tested, but no warranty is expressed or
  18. # implied.
  19. #
  20. # 13-Dec-2000: Updated to be used with the new bsddb3 package.
  21. # Added DBShelfCursor class.
  22. #
  23. #------------------------------------------------------------------------
  24. """Manage shelves of pickled objects using bsddb database files for the
  25. storage.
  26. """
  27. #------------------------------------------------------------------------
  28. import sys
  29. absolute_import = (sys.version_info[0] >= 3)
  30. if absolute_import :
  31. # Because this syntaxis is not valid before Python 2.5
  32. exec("from . import db")
  33. else :
  34. import db
  35. if sys.version_info[0] >= 3 :
  36. import cPickle # Will be converted to "pickle" by "2to3"
  37. else :
  38. if sys.version_info < (2, 6) :
  39. import cPickle
  40. else :
  41. # When we drop support for python 2.4
  42. # we could use: (in 2.5 we need a __future__ statement)
  43. #
  44. # with warnings.catch_warnings():
  45. # warnings.filterwarnings(...)
  46. # ...
  47. #
  48. # We can not use "with" as is, because it would be invalid syntax
  49. # in python 2.4 and (with no __future__) 2.5.
  50. # Here we simulate "with" following PEP 343 :
  51. import warnings
  52. w = warnings.catch_warnings()
  53. w.__enter__()
  54. try :
  55. warnings.filterwarnings('ignore',
  56. message='the cPickle module has been removed in Python 3.0',
  57. category=DeprecationWarning)
  58. import cPickle
  59. finally :
  60. w.__exit__()
  61. del w
  62. HIGHEST_PROTOCOL = cPickle.HIGHEST_PROTOCOL
  63. def _dumps(object, protocol):
  64. return cPickle.dumps(object, protocol=protocol)
  65. if sys.version_info < (2, 6) :
  66. from UserDict import DictMixin as MutableMapping
  67. else :
  68. import collections
  69. MutableMapping = collections.MutableMapping
  70. #------------------------------------------------------------------------
  71. def open(filename, flags=db.DB_CREATE, mode=0660, filetype=db.DB_HASH,
  72. dbenv=None, dbname=None):
  73. """
  74. A simple factory function for compatibility with the standard
  75. shleve.py module. It can be used like this, where key is a string
  76. and data is a pickleable object:
  77. from bsddb import dbshelve
  78. db = dbshelve.open(filename)
  79. db[key] = data
  80. db.close()
  81. """
  82. if type(flags) == type(''):
  83. sflag = flags
  84. if sflag == 'r':
  85. flags = db.DB_RDONLY
  86. elif sflag == 'rw':
  87. flags = 0
  88. elif sflag == 'w':
  89. flags = db.DB_CREATE
  90. elif sflag == 'c':
  91. flags = db.DB_CREATE
  92. elif sflag == 'n':
  93. flags = db.DB_TRUNCATE | db.DB_CREATE
  94. else:
  95. raise db.DBError, "flags should be one of 'r', 'w', 'c' or 'n' or use the bsddb.db.DB_* flags"
  96. d = DBShelf(dbenv)
  97. d.open(filename, dbname, filetype, flags, mode)
  98. return d
  99. #---------------------------------------------------------------------------
  100. class DBShelveError(db.DBError): pass
  101. class DBShelf(MutableMapping):
  102. """A shelf to hold pickled objects, built upon a bsddb DB object. It
  103. automatically pickles/unpickles data objects going to/from the DB.
  104. """
  105. def __init__(self, dbenv=None):
  106. self.db = db.DB(dbenv)
  107. self._closed = True
  108. if HIGHEST_PROTOCOL:
  109. self.protocol = HIGHEST_PROTOCOL
  110. else:
  111. self.protocol = 1
  112. def __del__(self):
  113. self.close()
  114. def __getattr__(self, name):
  115. """Many methods we can just pass through to the DB object.
  116. (See below)
  117. """
  118. return getattr(self.db, name)
  119. #-----------------------------------
  120. # Dictionary access methods
  121. def __len__(self):
  122. return len(self.db)
  123. def __getitem__(self, key):
  124. data = self.db[key]
  125. return cPickle.loads(data)
  126. def __setitem__(self, key, value):
  127. data = _dumps(value, self.protocol)
  128. self.db[key] = data
  129. def __delitem__(self, key):
  130. del self.db[key]
  131. def keys(self, txn=None):
  132. if txn is not None:
  133. return self.db.keys(txn)
  134. else:
  135. return self.db.keys()
  136. if sys.version_info >= (2, 6) :
  137. def __iter__(self) : # XXX: Load all keys in memory :-(
  138. for k in self.db.keys() :
  139. yield k
  140. # Do this when "DB" support iteration
  141. # Or is it enough to pass thru "getattr"?
  142. #
  143. # def __iter__(self) :
  144. # return self.db.__iter__()
  145. def open(self, *args, **kwargs):
  146. self.db.open(*args, **kwargs)
  147. self._closed = False
  148. def close(self, *args, **kwargs):
  149. self.db.close(*args, **kwargs)
  150. self._closed = True
  151. def __repr__(self):
  152. if self._closed:
  153. return '<DBShelf @ 0x%x - closed>' % (id(self))
  154. else:
  155. return repr(dict(self.iteritems()))
  156. def items(self, txn=None):
  157. if txn is not None:
  158. items = self.db.items(txn)
  159. else:
  160. items = self.db.items()
  161. newitems = []
  162. for k, v in items:
  163. newitems.append( (k, cPickle.loads(v)) )
  164. return newitems
  165. def values(self, txn=None):
  166. if txn is not None:
  167. values = self.db.values(txn)
  168. else:
  169. values = self.db.values()
  170. return map(cPickle.loads, values)
  171. #-----------------------------------
  172. # Other methods
  173. def __append(self, value, txn=None):
  174. data = _dumps(value, self.protocol)
  175. return self.db.append(data, txn)
  176. def append(self, value, txn=None):
  177. if self.get_type() == db.DB_RECNO:
  178. return self.__append(value, txn=txn)
  179. raise DBShelveError, "append() only supported when dbshelve opened with filetype=dbshelve.db.DB_RECNO"
  180. def associate(self, secondaryDB, callback, flags=0):
  181. def _shelf_callback(priKey, priData, realCallback=callback):
  182. # Safe in Python 2.x because expresion short circuit
  183. if sys.version_info[0] < 3 or isinstance(priData, bytes) :
  184. data = cPickle.loads(priData)
  185. else :
  186. data = cPickle.loads(bytes(priData, "iso8859-1")) # 8 bits
  187. return realCallback(priKey, data)
  188. return self.db.associate(secondaryDB, _shelf_callback, flags)
  189. #def get(self, key, default=None, txn=None, flags=0):
  190. def get(self, *args, **kw):
  191. # We do it with *args and **kw so if the default value wasn't
  192. # given nothing is passed to the extension module. That way
  193. # an exception can be raised if set_get_returns_none is turned
  194. # off.
  195. data = self.db.get(*args, **kw)
  196. try:
  197. return cPickle.loads(data)
  198. except (EOFError, TypeError, cPickle.UnpicklingError):
  199. return data # we may be getting the default value, or None,
  200. # so it doesn't need unpickled.
  201. def get_both(self, key, value, txn=None, flags=0):
  202. data = _dumps(value, self.protocol)
  203. data = self.db.get(key, data, txn, flags)
  204. return cPickle.loads(data)
  205. def cursor(self, txn=None, flags=0):
  206. c = DBShelfCursor(self.db.cursor(txn, flags))
  207. c.protocol = self.protocol
  208. return c
  209. def put(self, key, value, txn=None, flags=0):
  210. data = _dumps(value, self.protocol)
  211. return self.db.put(key, data, txn, flags)
  212. def join(self, cursorList, flags=0):
  213. raise NotImplementedError
  214. #----------------------------------------------
  215. # Methods allowed to pass-through to self.db
  216. #
  217. # close, delete, fd, get_byteswapped, get_type, has_key,
  218. # key_range, open, remove, rename, stat, sync,
  219. # upgrade, verify, and all set_* methods.
  220. #---------------------------------------------------------------------------
  221. class DBShelfCursor:
  222. """
  223. """
  224. def __init__(self, cursor):
  225. self.dbc = cursor
  226. def __del__(self):
  227. self.close()
  228. def __getattr__(self, name):
  229. """Some methods we can just pass through to the cursor object. (See below)"""
  230. return getattr(self.dbc, name)
  231. #----------------------------------------------
  232. def dup(self, flags=0):
  233. c = DBShelfCursor(self.dbc.dup(flags))
  234. c.protocol = self.protocol
  235. return c
  236. def put(self, key, value, flags=0):
  237. data = _dumps(value, self.protocol)
  238. return self.dbc.put(key, data, flags)
  239. def get(self, *args):
  240. count = len(args) # a method overloading hack
  241. method = getattr(self, 'get_%d' % count)
  242. method(*args)
  243. def get_1(self, flags):
  244. rec = self.dbc.get(flags)
  245. return self._extract(rec)
  246. def get_2(self, key, flags):
  247. rec = self.dbc.get(key, flags)
  248. return self._extract(rec)
  249. def get_3(self, key, value, flags):
  250. data = _dumps(value, self.protocol)
  251. rec = self.dbc.get(key, flags)
  252. return self._extract(rec)
  253. def current(self, flags=0): return self.get_1(flags|db.DB_CURRENT)
  254. def first(self, flags=0): return self.get_1(flags|db.DB_FIRST)
  255. def last(self, flags=0): return self.get_1(flags|db.DB_LAST)
  256. def next(self, flags=0): return self.get_1(flags|db.DB_NEXT)
  257. def prev(self, flags=0): return self.get_1(flags|db.DB_PREV)
  258. def consume(self, flags=0): return self.get_1(flags|db.DB_CONSUME)
  259. def next_dup(self, flags=0): return self.get_1(flags|db.DB_NEXT_DUP)
  260. def next_nodup(self, flags=0): return self.get_1(flags|db.DB_NEXT_NODUP)
  261. def prev_nodup(self, flags=0): return self.get_1(flags|db.DB_PREV_NODUP)
  262. def get_both(self, key, value, flags=0):
  263. data = _dumps(value, self.protocol)
  264. rec = self.dbc.get_both(key, flags)
  265. return self._extract(rec)
  266. def set(self, key, flags=0):
  267. rec = self.dbc.set(key, flags)
  268. return self._extract(rec)
  269. def set_range(self, key, flags=0):
  270. rec = self.dbc.set_range(key, flags)
  271. return self._extract(rec)
  272. def set_recno(self, recno, flags=0):
  273. rec = self.dbc.set_recno(recno, flags)
  274. return self._extract(rec)
  275. set_both = get_both
  276. def _extract(self, rec):
  277. if rec is None:
  278. return None
  279. else:
  280. key, data = rec
  281. # Safe in Python 2.x because expresion short circuit
  282. if sys.version_info[0] < 3 or isinstance(data, bytes) :
  283. return key, cPickle.loads(data)
  284. else :
  285. return key, cPickle.loads(bytes(data, "iso8859-1")) # 8 bits
  286. #----------------------------------------------
  287. # Methods allowed to pass-through to self.dbc
  288. #
  289. # close, count, delete, get_recno, join_item
  290. #---------------------------------------------------------------------------