src/engine/SCons/dblite.py

   1 # dblite.py module contributed by Ralf W. Grosse-Kunstleve.
   2 # Extended for Unicode by Steven Knight.
   3
   4 import cPickle
   5 import time
   6 import shutil
   7 import os
   8 import types
   9 import __builtin__
  10
  11 keep_all_files = 00000
  12 ignore_corrupt_dbfiles = 0
  13
  14 def corruption_warning(filename):
  15     print "Warning: Discarding corrupt database:", filename
  16
  17 if hasattr(types, 'UnicodeType'):
  18     def is_string(s):
  19         t = type(s)
  20         return t is types.StringType or t is types.UnicodeType
  21 else:
  22     def is_string(s):
  23         return type(s) is types.StringType
  24
  25 try:
  26     unicode('a')
  27 except NameError:
  28     def unicode(s): return s
  29
  30 dblite_suffix = '.dblite'
  31 tmp_suffix = '.tmp'
  32
  33 class dblite:
  34
  35   # Squirrel away references to the functions in various modules
  36   # that we'll use when our __del__() method calls our sync() method
  37   # during shutdown.  We might get destroyed when Python is in the midst
  38   # of tearing down the different modules we import in an essentially
  39   # arbitrary order, and some of the various modules's global attributes
  40   # may already be wiped out from under us.
  41   #
  42   # See the discussion at:
  43   #   http://mail.python.org/pipermail/python-bugs-list/2003-March/016877.html
  44
  45   _open = __builtin__.open
  46   _cPickle_dump = cPickle.dump
  47   _os_chmod = os.chmod
  48   try:
  49       _os_chown = os.chown
  50   except AttributeError:
  51       _os_chown = None
  52   _os_rename = os.rename
  53   _os_unlink = os.unlink
  54   _shutil_copyfile = shutil.copyfile
  55   _time_time = time.time
  56
  57   def __init__(self, file_base_name, flag, mode):
  58     assert flag in (None, "r", "w", "c", "n")
  59     if (flag is None): flag = "r"
  60     base, ext = os.path.splitext(file_base_name)
  61     if ext == dblite_suffix:
  62       # There's already a suffix on the file name, don't add one.
  63       self._file_name = file_base_name
  64       self._tmp_name = base + tmp_suffix
  65     else:
  66       self._file_name = file_base_name + dblite_suffix
  67       self._tmp_name = file_base_name + tmp_suffix
  68     self._flag = flag
  69     self._mode = mode
  70     self._dict = {}
  71     self._needs_sync = 00000
  72     if self._os_chown is not None and (os.geteuid()==0 or os.getuid()==0):
  73       # running as root; chown back to current owner/group when done
  74       try:
  75         statinfo = os.stat(self._file_name)
  76         self._chown_to = statinfo.st_uid
  77         self._chgrp_to = statinfo.st_gid
  78       except OSError, e:
  79         # db file doesn't exist yet.
  80         # Check os.environ for SUDO_UID, use if set
  81         self._chown_to = int(os.environ.get('SUDO_UID', -1))
  82         self._chgrp_to = int(os.environ.get('SUDO_GID', -1))
  83     else:
  84       self._chown_to = -1        # don't chown
  85       self._chgrp_to = -1        # don't chgrp
  86     if (self._flag == "n"):
  87       self._open(self._file_name, "wb", self._mode)
  88     else:
  89       try:
  90         f = self._open(self._file_name, "rb")
  91       except IOError, e:
  92         if (self._flag != "c"):
  93           raise e
  94         self._open(self._file_name, "wb", self._mode)
  95       else:
  96         p = f.read()
  97         if (len(p) > 0):
  98           try:
  99             self._dict = cPickle.loads(p)
 100           except (cPickle.UnpicklingError, EOFError):
 101             if (ignore_corrupt_dbfiles == 0): raise
 102             if (ignore_corrupt_dbfiles == 1):
 103               corruption_warning(self._file_name)
 104
 105   def __del__(self):
 106     if (self._needs_sync):
 107       self.sync()
 108
 109   def sync(self):
 110     self._check_writable()
 111     f = self._open(self._tmp_name, "wb", self._mode)
 112     self._cPickle_dump(self._dict, f, 1)
 113     f.close()
 114     # Windows doesn't allow renaming if the file exists, so unlink
 115     # it first, chmod'ing it to make sure we can do so.  On UNIX, we
 116     # may not be able to chmod the file if it's owned by someone else
 117     # (e.g. from a previous run as root).  We should still be able to
 118     # unlink() the file if the directory's writable, though, so ignore
 119     # any OSError exception  thrown by the chmod() call.
 120     try: self._os_chmod(self._file_name, 0777)
 121     except OSError: pass
 122     self._os_unlink(self._file_name)
 123     self._os_rename(self._tmp_name, self._file_name)
 124     if self._os_chown is not None and self._chown_to > 0: # don't chown to root or -1
 125       try:
 126         self._os_chown(self._file_name, self._chown_to, self._chgrp_to)
 127       except OSError:
 128         pass
 129     self._needs_sync = 00000
 130     if (keep_all_files):
 131       self._shutil_copyfile(
 132         self._file_name,
 133         self._file_name + "_" + str(int(self._time_time())))
 134
 135   def _check_writable(self):
 136     if (self._flag == "r"):
 137       raise IOError("Read-only database: %s" % self._file_name)
 138
 139   def __getitem__(self, key):
 140     return self._dict[key]
 141
 142   def __setitem__(self, key, value):
 143     self._check_writable()
 144     if (not is_string(key)):
 145       raise TypeError, "key `%s' must be a string but is %s" % (key, type(key))
 146     if (not is_string(value)):
 147       raise TypeError, "value `%s' must be a string but is %s" % (value, type(value))
 148     self._dict[key] = value
 149     self._needs_sync = 0001
 150
 151   def keys(self):
 152     return self._dict.keys()
 153
 154   def has_key(self, key):
 155     return key in self._dict
 156
 157   def __contains__(self, key):
 158     return key in self._dict
 159
 160   def iterkeys(self):
 161     return self._dict.iterkeys()
 162
 163   __iter__ = iterkeys
 164
 165   def __len__(self):
 166     return len(self._dict)
 167
 168 def open(file, flag=None, mode=0666):
 169   return dblite(file, flag, mode)
 170
 171 def _exercise():
 172   db = open("tmp", "n")
 173   assert len(db) == 0
 174   db["foo"] = "bar"
 175   assert db["foo"] == "bar"
 176   db[unicode("ufoo")] = unicode("ubar")
 177   assert db[unicode("ufoo")] == unicode("ubar")
 178   db.sync()
 179   db = open("tmp", "c")
 180   assert len(db) == 2, len(db)
 181   assert db["foo"] == "bar"
 182   db["bar"] = "foo"
 183   assert db["bar"] == "foo"
 184   db[unicode("ubar")] = unicode("ufoo")
 185   assert db[unicode("ubar")] == unicode("ufoo")
 186   db.sync()
 187   db = open("tmp", "r")
 188   assert len(db) == 4, len(db)
 189   assert db["foo"] == "bar"
 190   assert db["bar"] == "foo"
 191   assert db[unicode("ufoo")] == unicode("ubar")
 192   assert db[unicode("ubar")] == unicode("ufoo")
 193   try:
 194     db.sync()
 195   except IOError, e:
 196     assert str(e) == "Read-only database: tmp.dblite"
 197   else:
 198     raise RuntimeError, "IOError expected."
 199   db = open("tmp", "w")
 200   assert len(db) == 4
 201   db["ping"] = "pong"
 202   db.sync()
 203   try:
 204     db[(1,2)] = "tuple"
 205   except TypeError, e:
 206     assert str(e) == "key `(1, 2)' must be a string but is <type 'tuple'>", str(e)
 207   else:
 208     raise RuntimeError, "TypeError exception expected"
 209   try:
 210     db["list"] = [1,2]
 211   except TypeError, e:
 212     assert str(e) == "value `[1, 2]' must be a string but is <type 'list'>", str(e)
 213   else:
 214     raise RuntimeError, "TypeError exception expected"
 215   db = open("tmp", "r")
 216   assert len(db) == 5
 217   db = open("tmp", "n")
 218   assert len(db) == 0
 219   _open("tmp.dblite", "w")
 220   db = open("tmp", "r")
 221   _open("tmp.dblite", "w").write("x")
 222   try:
 223     db = open("tmp", "r")
 224   except cPickle.UnpicklingError:
 225     pass
 226   else:
 227     raise RuntimeError, "cPickle exception expected."
 228   global ignore_corrupt_dbfiles
 229   ignore_corrupt_dbfiles = 2
 230   db = open("tmp", "r")
 231   assert len(db) == 0
 232   os.unlink("tmp.dblite")
 233   try:
 234     db = open("tmp", "w")
 235   except IOError, e:
 236     assert str(e) == "[Errno 2] No such file or directory: 'tmp.dblite'", str(e)
 237   else:
 238     raise RuntimeError, "IOError expected."
 239   print "OK"
 240
 241 if (__name__ == "__main__"):
 242   _exercise()
 243
 244 # Local Variables:
 245 # tab-width:4
 246 # indent-tabs-mode:nil
 247 # End:
 248 # vim: set expandtab tabstop=4 shiftwidth=4: