Source Code for Module SCons.dblite

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