1 #!/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

    27 storage.

    28 """

    29 

    30 #------------------------------------------------------------------------

    31 

    32 import cPickle

    33 import sys

    34 

    35 import sys

    36 absolute_import = (sys.version_info[0] >= 3)

    37 if absolute_import :

    38     # Because this syntaxis is not valid before Python 2.5

    39     exec("from . import db")

    40 else :

    41     import db

    42 

    43 #At version 2.3 cPickle switched to using protocol instead of bin

    44 if sys.version_info[:3] >= (2, 3, 0):

    45     HIGHEST_PROTOCOL = cPickle.HIGHEST_PROTOCOL

    46 # In python 2.3.*, "cPickle.dumps" accepts no

    47 # named parameters. "pickle.dumps" accepts them,

    48 # so this seems a bug.

    49     if sys.version_info[:3] < (2, 4, 0):

    50         def _dumps(object, protocol):

    51             return cPickle.dumps(object, protocol)

    52     else :

    53         def _dumps(object, protocol):

    54             return cPickle.dumps(object, protocol=protocol)

    55 

    56 else:

    57     HIGHEST_PROTOCOL = None

    58     def _dumps(object, protocol):

    59         return cPickle.dumps(object, bin=protocol)

    60 

    61 

    62 if sys.version_info[0:2] <= (2, 5) :

    63     try:

    64         from UserDict import DictMixin

    65     except ImportError:

    66         # DictMixin is new in Python 2.3

    67         class DictMixin: pass

    68     MutableMapping = DictMixin

    69 else :

    70     import collections

    71     MutableMapping = collections.MutableMapping

    72 

    73 #------------------------------------------------------------------------

    74 

    75 

    76 def open(filename, flags=db.DB_CREATE, mode=0660, filetype=db.DB_HASH,

    77          dbenv=None, dbname=None):

    78     """

    79     A simple factory function for compatibility with the standard

    80     shleve.py module.  It can be used like this, where key is a string

    81     and data is a pickleable object:

    82 

    83         from bsddb import dbshelve

    84         db = dbshelve.open(filename)

    85 

    86         db[key] = data

    87 

    88         db.close()

    89     """

    90     if type(flags) == type(''):

    91         sflag = flags

    92         if sflag == 'r':

    93             flags = db.DB_RDONLY

    94         elif sflag == 'rw':

    95             flags = 0

    96         elif sflag == 'w':

    97             flags =  db.DB_CREATE

    98         elif sflag == 'c':

    99             flags =  db.DB_CREATE

   100         elif sflag == 'n':

   101             flags = db.DB_TRUNCATE | db.DB_CREATE

   102         else:

   103             raise db.DBError, "flags should be one of 'r', 'w', 'c' or 'n' or use the bsddb.db.DB_* flags"

   104 

   105     d = DBShelf(dbenv)

   106     d.open(filename, dbname, filetype, flags, mode)

   107     return d

   108 

   109 #---------------------------------------------------------------------------

   110 

   111 class DBShelveError(db.DBError): pass

   112 

   113 

   114 class DBShelf(MutableMapping):

   115     """A shelf to hold pickled objects, built upon a bsddb DB object.  It

   116     automatically pickles/unpickles data objects going to/from the DB.

   117     """

   118     def __init__(self, dbenv=None):

   119         self.db = db.DB(dbenv)

   120         self._closed = True

   121         if HIGHEST_PROTOCOL:

   122             self.protocol = HIGHEST_PROTOCOL

   123         else:

   124             self.protocol = 1

   125 

   126 

   127     def __del__(self):

   128         self.close()

   129 

   130 

   131     def __getattr__(self, name):

   132         """Many methods we can just pass through to the DB object.

   133         (See below)

   134         """

   135         return getattr(self.db, name)

   136 

   137 

   138     #-----------------------------------

   139     # Dictionary access methods

   140 

   141     def __len__(self):

   142         return len(self.db)

   143 

   144 

   145     def __getitem__(self, key):

   146         data = self.db[key]

   147         return cPickle.loads(data)

   148 

   149 

   150     def __setitem__(self, key, value):

   151         data = _dumps(value, self.protocol)

   152         self.db[key] = data

   153 

   154 

   155     def __delitem__(self, key):

   156         del self.db[key]

   157 

   158 

   159     def keys(self, txn=None):

   160         if txn != None:

   161             return self.db.keys(txn)

   162         else:

   163             return self.db.keys()

   164 

   165     if sys.version_info[0:2] >= (2, 6) :

   166         def __iter__(self) :

   167             return self.db.__iter__()

   168 

   169 

   170     def open(self, *args, **kwargs):

   171         self.db.open(*args, **kwargs)

   172         self._closed = False

   173 

   174 

   175     def close(self, *args, **kwargs):

   176         self.db.close(*args, **kwargs)

   177         self._closed = True

   178 

   179 

   180     def __repr__(self):

   181         if self._closed:

   182             return '<DBShelf @ 0x%x - closed>' % (id(self))

   183         else:

   184             return repr(dict(self.iteritems()))

   185 

   186 

   187     def items(self, txn=None):

   188         if txn != None:

   189             items = self.db.items(txn)

   190         else:

   191             items = self.db.items()

   192         newitems = []

   193 

   194         for k, v in items:

   195             newitems.append( (k, cPickle.loads(v)) )

   196         return newitems

   197 

   198     def values(self, txn=None):

   199         if txn != None:

   200             values = self.db.values(txn)

   201         else:

   202             values = self.db.values()

   203 

   204         return map(cPickle.loads, values)

   205 

   206     #-----------------------------------

   207     # Other methods

   208 

   209     def __append(self, value, txn=None):

   210         data = _dumps(value, self.protocol)

   211         return self.db.append(data, txn)

   212 

   213     def append(self, value, txn=None):

   214         if self.get_type() == db.DB_RECNO:

   215             return self.__append(value, txn=txn)

   216         raise DBShelveError, "append() only supported when dbshelve opened with filetype=dbshelve.db.DB_RECNO"

   217 

   218 

   219     def associate(self, secondaryDB, callback, flags=0):

   220         def _shelf_callback(priKey, priData, realCallback=callback):

   221             # Safe in Python 2.x because expresion short circuit

   222             if sys.version_info[0] < 3 or isinstance(priData, bytes) :

   223                 data = cPickle.loads(priData)

   224             else :

   225                 data = cPickle.loads(bytes(priData, "iso8859-1"))  # 8 bits

   226             return realCallback(priKey, data)

   227 

   228         return self.db.associate(secondaryDB, _shelf_callback, flags)

   229 

   230 

   231     #def get(self, key, default=None, txn=None, flags=0):

   232     def get(self, *args, **kw):

   233         # We do it with *args and **kw so if the default value wasn't

   234         # given nothing is passed to the extension module.  That way

   235         # an exception can be raised if set_get_returns_none is turned

   236         # off.

   237         data = apply(self.db.get, args, kw)

   238         try:

   239             return cPickle.loads(data)

   240         except (EOFError, TypeError, cPickle.UnpicklingError):

   241             return data  # we may be getting the default value, or None,

   242                          # so it doesn't need unpickled.

   243 

   244     def get_both(self, key, value, txn=None, flags=0):

   245         data = _dumps(value, self.protocol)

   246         data = self.db.get(key, data, txn, flags)

   247         return cPickle.loads(data)

   248 

   249 

   250     def cursor(self, txn=None, flags=0):

   251         c = DBShelfCursor(self.db.cursor(txn, flags))

   252         c.protocol = self.protocol

   253         return c

   254 

   255 

   256     def put(self, key, value, txn=None, flags=0):

   257         data = _dumps(value, self.protocol)

   258         return self.db.put(key, data, txn, flags)

   259 

   260 

   261     def join(self, cursorList, flags=0):

   262         raise NotImplementedError

   263 

   264 

   265     #----------------------------------------------

   266     # Methods allowed to pass-through to self.db

   267     #

   268     #    close,  delete, fd, get_byteswapped, get_type, has_key,

   269     #    key_range, open, remove, rename, stat, sync,

   270     #    upgrade, verify, and all set_* methods.

   271 

   272 

   273 #---------------------------------------------------------------------------

   274 

   275 class DBShelfCursor:

   276     """

   277     """

   278     def __init__(self, cursor):

   279         self.dbc = cursor

   280 

   281     def __del__(self):

   282         self.close()

   283 

   284 

   285     def __getattr__(self, name):

   286         """Some methods we can just pass through to the cursor object.  (See below)"""

   287         return getattr(self.dbc, name)

   288 

   289 

   290     #----------------------------------------------

   291 

   292     def dup(self, flags=0):

   293         c = DBShelfCursor(self.dbc.dup(flags))

   294         c.protocol = self.protocol

   295         return c

   296 

   297 

   298     def put(self, key, value, flags=0):

   299         data = _dumps(value, self.protocol)

   300         return self.dbc.put(key, data, flags)

   301 

   302 

   303     def get(self, *args):

   304         count = len(args)  # a method overloading hack

   305         method = getattr(self, 'get_%d' % count)

   306         apply(method, args)

   307 

   308     def get_1(self, flags):

   309         rec = self.dbc.get(flags)

   310         return self._extract(rec)

   311 

   312     def get_2(self, key, flags):

   313         rec = self.dbc.get(key, flags)

   314         return self._extract(rec)

   315 

   316     def get_3(self, key, value, flags):

   317         data = _dumps(value, self.protocol)

   318         rec = self.dbc.get(key, flags)

   319         return self._extract(rec)

   320 

   321 

   322     def current(self, flags=0): return self.get_1(flags|db.DB_CURRENT)

   323     def first(self, flags=0): return self.get_1(flags|db.DB_FIRST)

   324     def last(self, flags=0): return self.get_1(flags|db.DB_LAST)

   325     def next(self, flags=0): return self.get_1(flags|db.DB_NEXT)

   326     def prev(self, flags=0): return self.get_1(flags|db.DB_PREV)

   327     def consume(self, flags=0): return self.get_1(flags|db.DB_CONSUME)

   328     def next_dup(self, flags=0): return self.get_1(flags|db.DB_NEXT_DUP)

   329     def next_nodup(self, flags=0): return self.get_1(flags|db.DB_NEXT_NODUP)

   330     def prev_nodup(self, flags=0): return self.get_1(flags|db.DB_PREV_NODUP)

   331 

   332 

   333     def get_both(self, key, value, flags=0):

   334         data = _dumps(value, self.protocol)

   335         rec = self.dbc.get_both(key, flags)

   336         return self._extract(rec)

   337 

   338 

   339     def set(self, key, flags=0):

   340         rec = self.dbc.set(key, flags)

   341         return self._extract(rec)

   342 

   343     def set_range(self, key, flags=0):

   344         rec = self.dbc.set_range(key, flags)

   345         return self._extract(rec)

   346 

   347     def set_recno(self, recno, flags=0):

   348         rec = self.dbc.set_recno(recno, flags)

   349         return self._extract(rec)

   350 

   351     set_both = get_both

   352 

   353     def _extract(self, rec):

   354         if rec is None:

   355             return None

   356         else:

   357             key, data = rec

   358             # Safe in Python 2.x because expresion short circuit

   359             if sys.version_info[0] < 3 or isinstance(data, bytes) :

   360                 return key, cPickle.loads(data)

   361             else :

   362                 return key, cPickle.loads(bytes(data, "iso8859-1"))  # 8 bits

   363 

   364     #----------------------------------------------

   365     # Methods allowed to pass-through to self.dbc

   366     #

   367     # close, count, delete, get_recno, join_item

   368 

   369 

   370 #---------------------------------------------------------------------------