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