twisted/persisted/dirdbm.py

   1 # -*- test-case-name: twisted.test.test_dirdbm -*-
   2 #
   3 # Copyright (c) Twisted Matrix Laboratories.
   4 # See LICENSE for details.
   5
   6
   7
   8 """
   9 DBM-style interface to a directory.
  10
  11 Each key is stored as a single file.  This is not expected to be very fast or
  12 efficient, but it's good for easy debugging.
  13
  14 DirDBMs are *not* thread-safe, they should only be accessed by one thread at
  15 a time.
  16
  17 No files should be placed in the working directory of a DirDBM save those
  18 created by the DirDBM itself!
  19
  20 Maintainer: Itamar Shtull-Trauring
  21 """
  22
  23
  24 import os
  25 import types
  26 import base64
  27 import glob
  28
  29 try:
  30     import cPickle as pickle
  31 except ImportError:
  32     import pickle
  33
  34 try:
  35     _open
  36 except NameError:
  37     _open = open
  38
  39
  40 class DirDBM:
  41     """A directory with a DBM interface.
  42
  43     This class presents a hash-like interface to a directory of small,
  44     flat files. It can only use strings as keys or values.
  45     """
  46
  47     def __init__(self, name):
  48         """
  49         @type name: str
  50         @param name: Base path to use for the directory storage.
  51         """
  52         self.dname = os.path.abspath(name)
  53         if not os.path.isdir(self.dname):
  54             os.mkdir(self.dname)
  55         else:
  56             # Run recovery, in case we crashed. we delete all files ending
  57             # with ".new". Then we find all files who end with ".rpl". If a
  58             # corresponding file exists without ".rpl", we assume the write
  59             # failed and delete the ".rpl" file. If only a ".rpl" exist we
  60             # assume the program crashed right after deleting the old entry
  61             # but before renaming the replacement entry.
  62             #
  63             # NOTE: '.' is NOT in the base64 alphabet!
  64             for f in glob.glob(os.path.join(self.dname, "*.new")):
  65                 os.remove(f)
  66             replacements = glob.glob(os.path.join(self.dname, "*.rpl"))
  67             for f in replacements:
  68                 old = f[:-4]
  69                 if os.path.exists(old):
  70                     os.remove(f)
  71                 else:
  72                     os.rename(f, old)
  73
  74     def _encode(self, k):
  75         """Encode a key so it can be used as a filename.
  76         """
  77         # NOTE: '_' is NOT in the base64 alphabet!
  78         return base64.encodestring(k).replace('\n', '_').replace("/", "-")
  79
  80     def _decode(self, k):
  81         """Decode a filename to get the key.
  82         """
  83         return base64.decodestring(k.replace('_', '\n').replace("-", "/"))
  84
  85     def _readFile(self, path):
  86         """Read in the contents of a file.
  87
  88         Override in subclasses to e.g. provide transparently encrypted dirdbm.
  89         """
  90         f = _open(path, "rb")
  91         s = f.read()
  92         f.close()
  93         return s
  94
  95     def _writeFile(self, path, data):
  96         """Write data to a file.
  97
  98         Override in subclasses to e.g. provide transparently encrypted dirdbm.
  99         """
 100         f = _open(path, "wb")
 101         f.write(data)
 102         f.flush()
 103         f.close()
 104
 105     def __len__(self):
 106         """
 107         @return: The number of key/value pairs in this Shelf
 108         """
 109         return len(os.listdir(self.dname))
 110
 111     def __setitem__(self, k, v):
 112         """
 113         C{dirdbm[k] = v}
 114         Create or modify a textfile in this directory
 115
 116         @type k: str
 117         @param k: key to set
 118
 119         @type v: str
 120         @param v: value to associate with C{k}
 121         """
 122         assert type(k) == types.StringType, "DirDBM key must be a string"
 123         assert type(v) == types.StringType, "DirDBM value must be a string"
 124         k = self._encode(k)
 125
 126         # we create a new file with extension .new, write the data to it, and
 127         # if the write succeeds delete the old file and rename the new one.
 128         old = os.path.join(self.dname, k)
 129         if os.path.exists(old):
 130             new = old + ".rpl" # replacement entry
 131         else:
 132             new = old + ".new" # new entry
 133         try:
 134             self._writeFile(new, v)
 135         except:
 136             os.remove(new)
 137             raise
 138         else:
 139             if os.path.exists(old): os.remove(old)
 140             os.rename(new, old)
 141
 142     def __getitem__(self, k):
 143         """
 144         C{dirdbm[k]}
 145         Get the contents of a file in this directory as a string.
 146
 147         @type k: str
 148         @param k: key to lookup
 149
 150         @return: The value associated with C{k}
 151         @raise KeyError: Raised when there is no such key
 152         """
 153         assert type(k) == types.StringType, "DirDBM key must be a string"
 154         path = os.path.join(self.dname, self._encode(k))
 155         try:
 156             return self._readFile(path)
 157         except:
 158             raise KeyError, k
 159
 160     def __delitem__(self, k):
 161         """
 162         C{del dirdbm[foo]}
 163         Delete a file in this directory.
 164
 165         @type k: str
 166         @param k: key to delete
 167
 168         @raise KeyError: Raised when there is no such key
 169         """
 170         assert type(k) == types.StringType, "DirDBM key must be a string"
 171         k = self._encode(k)
 172         try:    os.remove(os.path.join(self.dname, k))
 173         except (OSError, IOError): raise KeyError(self._decode(k))
 174
 175     def keys(self):
 176         """
 177         @return: a C{list} of filenames (keys).
 178         """
 179         return map(self._decode, os.listdir(self.dname))
 180
 181     def values(self):
 182         """
 183         @return: a C{list} of file-contents (values).
 184         """
 185         vals = []
 186         keys = self.keys()
 187         for key in keys:
 188             vals.append(self[key])
 189         return vals
 190
 191     def items(self):
 192         """
 193         @return: a C{list} of 2-tuples containing key/value pairs.
 194         """
 195         items = []
 196         keys = self.keys()
 197         for key in keys:
 198             items.append((key, self[key]))
 199         return items
 200
 201     def has_key(self, key):
 202         """
 203         @type key: str
 204         @param key: The key to test
 205
 206         @return: A true value if this dirdbm has the specified key, a faluse
 207         value otherwise.
 208         """
 209         assert type(key) == types.StringType, "DirDBM key must be a string"
 210         key = self._encode(key)
 211         return os.path.isfile(os.path.join(self.dname, key))
 212
 213     def setdefault(self, key, value):
 214         """
 215         @type key: str
 216         @param key: The key to lookup
 217
 218         @param value: The value to associate with key if key is not already
 219         associated with a value.
 220         """
 221         if not self.has_key(key):
 222             self[key] = value
 223             return value
 224         return self[key]
 225
 226     def get(self, key, default = None):
 227         """
 228         @type key: str
 229         @param key: The key to lookup
 230
 231         @param default: The value to return if the given key does not exist
 232
 233         @return: The value associated with C{key} or C{default} if not
 234         C{self.has_key(key)}
 235         """
 236         if self.has_key(key):
 237             return self[key]
 238         else:
 239             return default
 240
 241     def __contains__(self, key):
 242         """
 243         C{key in dirdbm}
 244
 245         @type key: str
 246         @param key: The key to test
 247
 248         @return: A true value if C{self.has_key(key)}, a false value otherwise.
 249         """
 250         assert type(key) == types.StringType, "DirDBM key must be a string"
 251         key = self._encode(key)
 252         return os.path.isfile(os.path.join(self.dname, key))
 253
 254     def update(self, dict):
 255         """
 256         Add all the key/value pairs in C{dict} to this dirdbm.  Any conflicting
 257         keys will be overwritten with the values from C{dict}.
 258
 259         @type dict: mapping
 260         @param dict: A mapping of key/value pairs to add to this dirdbm.
 261         """
 262         for key, val in dict.items():
 263             self[key]=val
 264
 265     def copyTo(self, path):
 266         """
 267         Copy the contents of this dirdbm to the dirdbm at C{path}.
 268
 269         @type path: C{str}
 270         @param path: The path of the dirdbm to copy to.  If a dirdbm
 271         exists at the destination path, it is cleared first.
 272
 273         @rtype: C{DirDBM}
 274         @return: The dirdbm this dirdbm was copied to.
 275         """
 276         path = os.path.abspath(path)
 277         assert path != self.dname
 278
 279         d = self.__class__(path)
 280         d.clear()
 281         for k in self.keys():
 282             d[k] = self[k]
 283         return d
 284
 285     def clear(self):
 286         """
 287         Delete all key/value pairs in this dirdbm.
 288         """
 289         for k in self.keys():
 290             del self[k]
 291
 292     def close(self):
 293         """
 294         Close this dbm: no-op, for dbm-style interface compliance.
 295         """
 296
 297     def getModificationTime(self, key):
 298         """
 299         Returns modification time of an entry.
 300
 301         @return: Last modification date (seconds since epoch) of entry C{key}
 302         @raise KeyError: Raised when there is no such key
 303         """
 304         assert type(key) == types.StringType, "DirDBM key must be a string"
 305         path = os.path.join(self.dname, self._encode(key))
 306         if os.path.isfile(path):
 307             return os.path.getmtime(path)
 308         else:
 309             raise KeyError, key
 310
 311
 312 class Shelf(DirDBM):
 313     """A directory with a DBM shelf interface.
 314
 315     This class presents a hash-like interface to a directory of small,
 316     flat files. Keys must be strings, but values can be any given object.
 317     """
 318
 319     def __setitem__(self, k, v):
 320         """
 321         C{shelf[foo] = bar}
 322         Create or modify a textfile in this directory.
 323
 324         @type k: str
 325         @param k: The key to set
 326
 327         @param v: The value to associate with C{key}
 328         """
 329         v = pickle.dumps(v)
 330         DirDBM.__setitem__(self, k, v)
 331
 332     def __getitem__(self, k):
 333         """
 334         C{dirdbm[foo]}
 335         Get and unpickle the contents of a file in this directory.
 336
 337         @type k: str
 338         @param k: The key to lookup
 339
 340         @return: The value associated with the given key
 341         @raise KeyError: Raised if the given key does not exist
 342         """
 343         return pickle.loads(DirDBM.__getitem__(self, k))
 344
 345
 346 def open(file, flag = None, mode = None):
 347     """
 348     This is for 'anydbm' compatibility.
 349
 350     @param file: The parameter to pass to the DirDBM constructor.
 351
 352     @param flag: ignored
 353     @param mode: ignored
 354     """
 355     return DirDBM(file)
 356
 357
 358 __all__ = ["open", "DirDBM", "Shelf"]