pygr/dbfile.py

   1
   2 import shelve, anydbm, sys, UserDict
   3 import logger
   4
   5 class WrongFormatError(IOError):
   6     'attempted to open db with the wrong format e.g. btree vs. hash'
   7     pass
   8 class NoSuchFileError(IOError):
   9     'file does not exist!'
  10     pass
  11 class PermissionsError(IOError):
  12     'inadequate permissions for requested file'
  13     pass
  14 class ReadOnlyError(PermissionsError):
  15     'attempted to open a file for writing, but no write permission'
  16     pass
  17
  18 def open_anydbm(*args, **kwargs):
  19     'trap anydbm.error message and transform to our consistent exception types'
  20     try:
  21         return anydbm.open(*args, **kwargs)
  22     except anydbm.error, e:
  23         msg = str(e)
  24         if msg.endswith('new db'):
  25             raise NoSuchFileError(msg)
  26         elif msg.startswith('db type'):
  27             raise WrongFormatError(msg)
  28         raise
  29
  30 try: # detect whether bsddb module available and working...
  31     import bsddb
  32     try:
  33         bsddb.db
  34     except AttributeError:
  35         raise ImportError
  36 except ImportError:
  37     bsddb = None
  38
  39 def open_bsddb(filename, flag='r', useHash=False, mode=0666):
  40     """open bsddb index instead of hash by default.
  41     useHash=True forces it to use anydbm default (i.e. hash) instead.
  42     Also gives more meaningful error messages."""
  43     try: # 1ST OPEN AS BTREE
  44         if useHash: # FORCE IT TO USE HASH INSTEAD OF BTREE
  45             return open_anydbm(filename, flag)
  46         else:
  47             return bsddb.btopen(filename, flag, mode)
  48     except bsddb.db.DBAccessError: # HMM, BLOCKED BY PERMISSIONS
  49         if flag=='c' or flag=='w': # TRY OPENING READ-ONLY
  50             try:
  51                 ifile = file(filename)
  52             except IOError: # HMM, NOT EVEN READABLE. RAISE GENERIC PERMISSIONS ERROR
  53                 raise PermissionsError('insufficient permissions to open file: '
  54                                        +filename)
  55             ifile.close() # OK, WE CAN READ THE FILE, SO RAISE EXCEPTION WITH
  56             raise ReadOnlyError('file is read-only: '+filename) # VERY SPECIFIC MEANING!
  57         else: # r OR n FLAG: JUST RAISE EXCEPTION
  58             raise PermissionsError('insufficient permissions to open file: '
  59                                    +filename)
  60     except bsddb.db.DBNoSuchFileError:
  61         raise NoSuchFileError('no file named: '+filename)
  62     except bsddb.db.DBInvalidArgError: # NOT A BTREE FILE...
  63         try:
  64             if useHash: # NO POINT IN TRYING HASH YET AGAIN...
  65                 raise bsddb.db.DBInvalidArgError
  66             # fallback to using default: hash file
  67             return open_anydbm(filename, flag)
  68         except bsddb.db.DBInvalidArgError:
  69             raise WrongFormatError('file does not match expected shelve format: '+filename)
  70
  71 def open_index(filename, flag='r', useHash=False, mode=0666):
  72     if bsddb is None:
  73         d = open_anydbm(filename, flag)
  74         if not useHash:
  75             logger.warn('Falling back to hash index: unable to import bsddb')
  76         return d
  77     return open_bsddb(filename, flag, useHash, mode)
  78
  79 def iter_gdbm(db):
  80     'iterator for gdbm objects'
  81     k = db.firstkey()
  82     while k is not None:
  83         yield k
  84         k = db.nextkey(k)
  85
  86 class _ClosedDict(UserDict.DictMixin):
  87     """This dummy class exists solely to raise a clear error msg if accessed.
  88     Copied from the Python 2.6 shelve.py """
  89     def closed(self, *args):
  90         raise ValueError('invalid operation on closed shelf')
  91     __getitem__ = __setitem__ = __delitem__ = keys = closed
  92
  93     def __repr__(self):
  94         return '<Closed Dictionary>'
  95
  96
  97
  98 class BetterShelf(shelve.Shelf):
  99     """Shelf subclass that fixes its horrible iter implementation.
 100     """
 101     def __iter__(self):
 102         'avoid using iter provided by shelve/DictMixin, which loads all keys!'
 103         try:
 104             return iter(self.dict)
 105         except TypeError: # gdbm lacks __iter__ method, so try iter_gdbm()
 106             exc_type, exc_value, exc_traceback = sys.exc_info()
 107             try:
 108                 self.dict.firstkey
 109             except AttributeError: # evidently not a gdbm dict
 110                 raise exc_value, None, exc_traceback # re-raise original error
 111             else: # iterate using gdbm-specific method
 112                 return iter_gdbm(self.dict)
 113
 114     if sys.version_info < (2, 6): # Python finally added good err msg in 2.6
 115         def close(self):
 116             if isinstance(self.dict, _ClosedDict):
 117                 return # if already closed, nothing further to do...
 118             shelve.Shelf.close(self) # close Shelf as usual
 119             self.dict = _ClosedDict() # raises sensible error msg if accessed
 120
 121 def shelve_open(filename, flag='c', protocol=None, writeback=False,
 122                 useHash=True, mode=0666, *args, **kwargs):
 123     """improved implementation of shelve.open() that won't generate
 124 bogus __del__ warning messages like Python's version does."""
 125     d = open_index(filename, flag, useHash, mode) # construct Shelf only if OK
 126     return BetterShelf(d, protocol, writeback, *args, **kwargs)