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