1 """Manage shelves of pickled objects.

     2 

     3 A "shelf" is a persistent, dictionary-like object.  The difference

     4 with dbm databases is that the values (not the keys!) in a shelf can

     5 be essentially arbitrary Python objects -- anything that the "pickle"

     6 module can handle.  This includes most class instances, recursive data

     7 types, and objects containing lots of shared sub-objects.  The keys

     8 are ordinary strings.

     9 

    10 To summarize the interface (key is a string, data is an arbitrary

    11 object):

    12 

    13         import shelve

    14         d = shelve.open(filename) # open, with (g)dbm filename -- no suffix

    15 

    16         d[key] = data   # store data at key (overwrites old data if

    17                         # using an existing key)

    18         data = d[key]   # retrieve a COPY of the data at key (raise

    19                         # KeyError if no such key) -- NOTE that this

    20                         # access returns a *copy* of the entry!

    21         del d[key]      # delete data stored at key (raises KeyError

    22                         # if no such key)

    23         flag = d.has_key(key)   # true if the key exists; same as "key in d"

    24         list = d.keys() # a list of all existing keys (slow!)

    25 

    26         d.close()       # close it

    27 

    28 Dependent on the implementation, closing a persistent dictionary may

    29 or may not be necessary to flush changes to disk.

    30 

    31 Normally, d[key] returns a COPY of the entry.  This needs care when

    32 mutable entries are mutated: for example, if d[key] is a list,

    33         d[key].append(anitem)

    34 does NOT modify the entry d[key] itself, as stored in the persistent

    35 mapping -- it only modifies the copy, which is then immediately

    36 discarded, so that the append has NO effect whatsoever.  To append an

    37 item to d[key] in a way that will affect the persistent mapping, use:

    38         data = d[key]

    39         data.append(anitem)

    40         d[key] = data

    41 

    42 To avoid the problem with mutable entries, you may pass the keyword

    43 argument writeback=True in the call to shelve.open.  When you use:

    44         d = shelve.open(filename, writeback=True)

    45 then d keeps a cache of all entries you access, and writes them all back

    46 to the persistent mapping when you call d.close().  This ensures that

    47 such usage as d[key].append(anitem) works as intended.

    48 

    49 However, using keyword argument writeback=True may consume vast amount

    50 of memory for the cache, and it may make d.close() very slow, if you

    51 access many of d's entries after opening it in this way: d has no way to

    52 check which of the entries you access are mutable and/or which ones you

    53 actually mutate, so it must cache, and write back at close, all of the

    54 entries that you access.  You can call d.sync() to write back all the

    55 entries in the cache, and empty the cache (d.sync() also synchronizes

    56 the persistent dictionary on disk, if feasible).

    57 """

    58 

    59 # Try using cPickle and cStringIO if available.

    60 

    61 try:

    62     from cPickle import Pickler, Unpickler

    63 except ImportError:

    64     from pickle import Pickler, Unpickler

    65 

    66 try:

    67     from cStringIO import StringIO

    68 except ImportError:

    69     from StringIO import StringIO

    70 

    71 import UserDict

    72 import warnings

    73 

    74 __all__ = ["Shelf","BsdDbShelf","DbfilenameShelf","open"]

    75 

    76 class Shelf(UserDict.DictMixin):

    77     """Base class for shelf implementations.

    78 

    79     This is initialized with a dictionary-like object.

    80     See the module's __doc__ string for an overview of the interface.

    81     """

    82 

    83     def __init__(self, dict, protocol=None, writeback=False):

    84         self.dict = dict

    85         if protocol is None:

    86             protocol = 0

    87         self._protocol = protocol

    88         self.writeback = writeback

    89         self.cache = {}

    90 

    91     def keys(self):

    92         return self.dict.keys()

    93 

    94     def __len__(self):

    95         return len(self.dict)

    96 

    97     def has_key(self, key):

    98         return self.dict.has_key(key)

    99 

   100     def __contains__(self, key):

   101         return self.dict.has_key(key)

   102 

   103     def get(self, key, default=None):

   104         if self.dict.has_key(key):

   105             return self[key]

   106         return default

   107 

   108     def __getitem__(self, key):

   109         try:

   110             value = self.cache[key]

   111         except KeyError:

   112             f = StringIO(self.dict[key])

   113             value = Unpickler(f).load()

   114             if self.writeback:

   115                 self.cache[key] = value

   116         return value

   117 

   118     def __setitem__(self, key, value):

   119         if self.writeback:

   120             self.cache[key] = value

   121         f = StringIO()

   122         p = Pickler(f, self._protocol)

   123         p.dump(value)

   124         self.dict[key] = f.getvalue()

   125 

   126     def __delitem__(self, key):

   127         del self.dict[key]

   128         try:

   129             del self.cache[key]

   130         except KeyError:

   131             pass

   132 

   133     def close(self):

   134         self.sync()

   135         try:

   136             self.dict.close()

   137         except AttributeError:

   138             pass

   139         self.dict = 0

   140 

   141     def __del__(self):

   142         if not hasattr(self, 'writeback'):

   143             # __init__ didn't succeed, so don't bother closing

   144             return

   145         self.close()

   146 

   147     def sync(self):

   148         if self.writeback and self.cache:

   149             self.writeback = False

   150             for key, entry in self.cache.iteritems():

   151                 self[key] = entry

   152             self.writeback = True

   153             self.cache = {}

   154         if hasattr(self.dict, 'sync'):

   155             self.dict.sync()

   156 

   157 

   158 class BsdDbShelf(Shelf):

   159     """Shelf implementation using the "BSD" db interface.

   160 

   161     This adds methods first(), next(), previous(), last() and

   162     set_location() that have no counterpart in [g]dbm databases.

   163 

   164     The actual database must be opened using one of the "bsddb"

   165     modules "open" routines (i.e. bsddb.hashopen, bsddb.btopen or

   166     bsddb.rnopen) and passed to the constructor.

   167 

   168     See the module's __doc__ string for an overview of the interface.

   169     """

   170 

   171     def __init__(self, dict, protocol=None, writeback=False):

   172         Shelf.__init__(self, dict, protocol, writeback)

   173 

   174     def set_location(self, key):

   175         (key, value) = self.dict.set_location(key)

   176         f = StringIO(value)

   177         return (key, Unpickler(f).load())

   178 

   179     def next(self):

   180         (key, value) = self.dict.next()

   181         f = StringIO(value)

   182         return (key, Unpickler(f).load())

   183 

   184     def previous(self):

   185         (key, value) = self.dict.previous()

   186         f = StringIO(value)

   187         return (key, Unpickler(f).load())

   188 

   189     def first(self):

   190         (key, value) = self.dict.first()

   191         f = StringIO(value)

   192         return (key, Unpickler(f).load())

   193 

   194     def last(self):

   195         (key, value) = self.dict.last()

   196         f = StringIO(value)

   197         return (key, Unpickler(f).load())

   198 

   199 

   200 class DbfilenameShelf(Shelf):

   201     """Shelf implementation using the "anydbm" generic dbm interface.

   202 

   203     This is initialized with the filename for the dbm database.

   204     See the module's __doc__ string for an overview of the interface.

   205     """

   206 

   207     def __init__(self, filename, flag='c', protocol=None, writeback=False):

   208         import anydbm

   209         Shelf.__init__(self, anydbm.open(filename, flag), protocol, writeback)

   210 

   211 

   212 def open(filename, flag='c', protocol=None, writeback=False):

   213     """Open a persistent dictionary for reading and writing.

   214 

   215     The filename parameter is the base filename for the underlying

   216     database.  As a side-effect, an extension may be added to the

   217     filename and more than one file may be created.  The optional flag

   218     parameter has the same interpretation as the flag parameter of

   219     anydbm.open(). The optional protocol parameter specifies the

   220     version of the pickle protocol (0, 1, or 2).

   221 

   222     See the module's __doc__ string for an overview of the interface.

   223     """

   224 

   225     return DbfilenameShelf(filename, flag, protocol, writeback)