Source Code for Module SCons.dblite

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