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
25"""Manage shelves of pickled objects using bsddb database files for the
26storage.
27"""
28
29#------------------------------------------------------------------------
30
31import sys
32absolute_import = (sys.version_info[0] >= 3)
33if absolute_import :
34    # Because this syntaxis is not valid before Python 2.5
35    exec("from . import db")
36else :
37    import db
38
39if sys.version_info[0] >= 3 :
40    import cPickle  # Will be converted to "pickle" by "2to3"
41else :
42    if sys.version_info < (2, 6) :
43        import cPickle
44    else :
45        # When we drop support for python 2.4
46        # we could use: (in 2.5 we need a __future__ statement)
47        #
48        #    with warnings.catch_warnings():
49        #        warnings.filterwarnings(...)
50        #        ...
51        #
52        # We can not use "with" as is, because it would be invalid syntax
53        # in python 2.4 and (with no __future__) 2.5.
54        # Here we simulate "with" following PEP 343 :
55        import warnings
56        w = warnings.catch_warnings()
57        w.__enter__()
58        try :
59            warnings.filterwarnings('ignore',
60                message='the cPickle module has been removed in Python 3.0',
61                category=DeprecationWarning)
62            import cPickle
63        finally :
64            w.__exit__()
65        del w
66
67HIGHEST_PROTOCOL = cPickle.HIGHEST_PROTOCOL
68def _dumps(object, protocol):
69    return cPickle.dumps(object, protocol=protocol)
70
71if sys.version_info < (2, 6) :
72    from UserDict import DictMixin as MutableMapping
73else :
74    import collections
75    MutableMapping = collections.MutableMapping
76
77#------------------------------------------------------------------------
78
79
80def open(filename, flags=db.DB_CREATE, mode=0660, filetype=db.DB_HASH,
81         dbenv=None, dbname=None):
82    """
83    A simple factory function for compatibility with the standard
84    shleve.py module.  It can be used like this, where key is a string
85    and data is a pickleable object:
86
87        from bsddb import dbshelve
88        db = dbshelve.open(filename)
89
90        db[key] = data
91
92        db.close()
93    """
94    if type(flags) == type(''):
95        sflag = flags
96        if sflag == 'r':
97            flags = db.DB_RDONLY
98        elif sflag == 'rw':
99            flags = 0
100        elif sflag == 'w':
101            flags =  db.DB_CREATE
102        elif sflag == 'c':
103            flags =  db.DB_CREATE
104        elif sflag == 'n':
105            flags = db.DB_TRUNCATE | db.DB_CREATE
106        else:
107            raise db.DBError, "flags should be one of 'r', 'w', 'c' or 'n' or use the bsddb.db.DB_* flags"
108
109    d = DBShelf(dbenv)
110    d.open(filename, dbname, filetype, flags, mode)
111    return d
112
113#---------------------------------------------------------------------------
114
115class DBShelveError(db.DBError): pass
116
117
118class DBShelf(MutableMapping):
119    """A shelf to hold pickled objects, built upon a bsddb DB object.  It
120    automatically pickles/unpickles data objects going to/from the DB.
121    """
122    def __init__(self, dbenv=None):
123        self.db = db.DB(dbenv)
124        self._closed = True
125        if HIGHEST_PROTOCOL:
126            self.protocol = HIGHEST_PROTOCOL
127        else:
128            self.protocol = 1
129
130
131    def __del__(self):
132        self.close()
133
134
135    def __getattr__(self, name):
136        """Many methods we can just pass through to the DB object.
137        (See below)
138        """
139        return getattr(self.db, name)
140
141
142    #-----------------------------------
143    # Dictionary access methods
144
145    def __len__(self):
146        return len(self.db)
147
148
149    def __getitem__(self, key):
150        data = self.db[key]
151        return cPickle.loads(data)
152
153
154    def __setitem__(self, key, value):
155        data = _dumps(value, self.protocol)
156        self.db[key] = data
157
158
159    def __delitem__(self, key):
160        del self.db[key]
161
162
163    def keys(self, txn=None):
164        if txn is not None:
165            return self.db.keys(txn)
166        else:
167            return self.db.keys()
168
169    if sys.version_info >= (2, 6) :
170        def __iter__(self) :  # XXX: Load all keys in memory :-(
171            for k in self.db.keys() :
172                yield k
173
174        # Do this when "DB"  support iteration
175        # Or is it enough to pass thru "getattr"?
176        #
177        # def __iter__(self) :
178        #    return self.db.__iter__()
179
180
181    def open(self, *args, **kwargs):
182        self.db.open(*args, **kwargs)
183        self._closed = False
184
185
186    def close(self, *args, **kwargs):
187        self.db.close(*args, **kwargs)
188        self._closed = True
189
190
191    def __repr__(self):
192        if self._closed:
193            return '<DBShelf @ 0x%x - closed>' % (id(self))
194        else:
195            return repr(dict(self.iteritems()))
196
197
198    def items(self, txn=None):
199        if txn is not None:
200            items = self.db.items(txn)
201        else:
202            items = self.db.items()
203        newitems = []
204
205        for k, v in items:
206            newitems.append( (k, cPickle.loads(v)) )
207        return newitems
208
209    def values(self, txn=None):
210        if txn is not None:
211            values = self.db.values(txn)
212        else:
213            values = self.db.values()
214
215        return map(cPickle.loads, values)
216
217    #-----------------------------------
218    # Other methods
219
220    def __append(self, value, txn=None):
221        data = _dumps(value, self.protocol)
222        return self.db.append(data, txn)
223
224    def append(self, value, txn=None):
225        if self.get_type() == db.DB_RECNO:
226            return self.__append(value, txn=txn)
227        raise DBShelveError, "append() only supported when dbshelve opened with filetype=dbshelve.db.DB_RECNO"
228
229
230    def associate(self, secondaryDB, callback, flags=0):
231        def _shelf_callback(priKey, priData, realCallback=callback):
232            # Safe in Python 2.x because expresion short circuit
233            if sys.version_info[0] < 3 or isinstance(priData, bytes) :
234                data = cPickle.loads(priData)
235            else :
236                data = cPickle.loads(bytes(priData, "iso8859-1"))  # 8 bits
237            return realCallback(priKey, data)
238
239        return self.db.associate(secondaryDB, _shelf_callback, flags)
240
241
242    #def get(self, key, default=None, txn=None, flags=0):
243    def get(self, *args, **kw):
244        # We do it with *args and **kw so if the default value wasn't
245        # given nothing is passed to the extension module.  That way
246        # an exception can be raised if set_get_returns_none is turned
247        # off.
248        data = self.db.get(*args, **kw)
249        try:
250            return cPickle.loads(data)
251        except (EOFError, TypeError, cPickle.UnpicklingError):
252            return data  # we may be getting the default value, or None,
253                         # so it doesn't need unpickled.
254
255    def get_both(self, key, value, txn=None, flags=0):
256        data = _dumps(value, self.protocol)
257        data = self.db.get(key, data, txn, flags)
258        return cPickle.loads(data)
259
260
261    def cursor(self, txn=None, flags=0):
262        c = DBShelfCursor(self.db.cursor(txn, flags))
263        c.protocol = self.protocol
264        return c
265
266
267    def put(self, key, value, txn=None, flags=0):
268        data = _dumps(value, self.protocol)
269        return self.db.put(key, data, txn, flags)
270
271
272    def join(self, cursorList, flags=0):
273        raise NotImplementedError
274
275
276    #----------------------------------------------
277    # Methods allowed to pass-through to self.db
278    #
279    #    close,  delete, fd, get_byteswapped, get_type, has_key,
280    #    key_range, open, remove, rename, stat, sync,
281    #    upgrade, verify, and all set_* methods.
282
283
284#---------------------------------------------------------------------------
285
286class DBShelfCursor:
287    """
288    """
289    def __init__(self, cursor):
290        self.dbc = cursor
291
292    def __del__(self):
293        self.close()
294
295
296    def __getattr__(self, name):
297        """Some methods we can just pass through to the cursor object.  (See below)"""
298        return getattr(self.dbc, name)
299
300
301    #----------------------------------------------
302
303    def dup(self, flags=0):
304        c = DBShelfCursor(self.dbc.dup(flags))
305        c.protocol = self.protocol
306        return c
307
308
309    def put(self, key, value, flags=0):
310        data = _dumps(value, self.protocol)
311        return self.dbc.put(key, data, flags)
312
313
314    def get(self, *args):
315        count = len(args)  # a method overloading hack
316        method = getattr(self, 'get_%d' % count)
317        method(*args)
318
319    def get_1(self, flags):
320        rec = self.dbc.get(flags)
321        return self._extract(rec)
322
323    def get_2(self, key, flags):
324        rec = self.dbc.get(key, flags)
325        return self._extract(rec)
326
327    def get_3(self, key, value, flags):
328        data = _dumps(value, self.protocol)
329        rec = self.dbc.get(key, flags)
330        return self._extract(rec)
331
332
333    def current(self, flags=0): return self.get_1(flags|db.DB_CURRENT)
334    def first(self, flags=0): return self.get_1(flags|db.DB_FIRST)
335    def last(self, flags=0): return self.get_1(flags|db.DB_LAST)
336    def next(self, flags=0): return self.get_1(flags|db.DB_NEXT)
337    def prev(self, flags=0): return self.get_1(flags|db.DB_PREV)
338    def consume(self, flags=0): return self.get_1(flags|db.DB_CONSUME)
339    def next_dup(self, flags=0): return self.get_1(flags|db.DB_NEXT_DUP)
340    def next_nodup(self, flags=0): return self.get_1(flags|db.DB_NEXT_NODUP)
341    def prev_nodup(self, flags=0): return self.get_1(flags|db.DB_PREV_NODUP)
342
343
344    def get_both(self, key, value, flags=0):
345        data = _dumps(value, self.protocol)
346        rec = self.dbc.get_both(key, flags)
347        return self._extract(rec)
348
349
350    def set(self, key, flags=0):
351        rec = self.dbc.set(key, flags)
352        return self._extract(rec)
353
354    def set_range(self, key, flags=0):
355        rec = self.dbc.set_range(key, flags)
356        return self._extract(rec)
357
358    def set_recno(self, recno, flags=0):
359        rec = self.dbc.set_recno(recno, flags)
360        return self._extract(rec)
361
362    set_both = get_both
363
364    def _extract(self, rec):
365        if rec is None:
366            return None
367        else:
368            key, data = rec
369            # Safe in Python 2.x because expresion short circuit
370            if sys.version_info[0] < 3 or isinstance(data, bytes) :
371                return key, cPickle.loads(data)
372            else :
373                return key, cPickle.loads(bytes(data, "iso8859-1"))  # 8 bits
374
375    #----------------------------------------------
376    # Methods allowed to pass-through to self.dbc
377    #
378    # close, count, delete, get_recno, join_item
379
380
381#---------------------------------------------------------------------------
382