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              self._open(self._file_name, "wb", self._mode) 
112          else: 
113              try: 
114                  f = self._open(self._file_name, "rb") 
115              except IOError as e: 
116                  if self._flag != "c": 
117                      raise e 
118                  self._open(self._file_name, "wb", self._mode) 
119              else: 
120                  p = f.read() 
121                  if len(p) > 0: 
122                      try: 
123                          if bytes is not str: 
124                              self._dict = pickle.loads(p, encoding='bytes') 
125                          else: 
126                              self._dict = pickle.loads(p) 
127                      except (pickle.UnpicklingError, EOFError, KeyError): 
128                          # Note how we catch KeyErrors too here, which might happen 
129                          # when we don't have cPickle available (default pickle 
130                          # throws it). 
131                          if (ignore_corrupt_dbfiles == 0): raise 
132                          if (ignore_corrupt_dbfiles == 1): 
133                              corruption_warning(self._file_name) 
134   
135 -    def close(self): 
136          if self._needs_sync: 
137              self.sync() 
138   
139 -    def __del__(self): 
140          self.close() 
141   
142 -    def sync(self): 
143          self._check_writable() 
144          f = self._open(self._tmp_name, "wb", self._mode) 
145          self._pickle_dump(self._dict, f, self._pickle_protocol) 
146          f.close() 
147   
148          # Windows doesn't allow renaming if the file exists, so unlink 
149          # it first, chmod'ing it to make sure we can do so.  On UNIX, we 
150          # may not be able to chmod the file if it's owned by someone else 
151          # (e.g. from a previous run as root).  We should still be able to 
152          # unlink() the file if the directory's writable, though, so ignore 
153          # any OSError exception  thrown by the chmod() call. 
154          try: 
155              self._os_chmod(self._file_name, 0o777) 
156          except OSError: 
157              pass 
158          self._os_unlink(self._file_name) 
159          self._os_rename(self._tmp_name, self._file_name) 
160          if self._os_chown is not None and self._chown_to > 0:  # don't chown to root or -1 
161              try: 
162                  self._os_chown(self._file_name, self._chown_to, self._chgrp_to) 
163              except OSError: 
164                  pass 
165          self._needs_sync = 00000 
166          if (keep_all_files): 
167              self._shutil_copyfile( 
168                  self._file_name, 
169                  self._file_name + "_" + str(int(self._time_time()))) 
170   
171 -    def _check_writable(self): 
172          if (self._flag == "r"): 
173              raise IOError("Read-only database: %s" % self._file_name) 
174   
175 -    def __getitem__(self, key): 
176          return self._dict[key] 
177   
178 -    def __setitem__(self, key, value): 
179          self._check_writable() 
180          if (not is_string(key)): 
181              raise TypeError("key `%s' must be a string but is %s" % (key, type(key))) 
182          if (not is_bytes(value)): 
183              raise TypeError("value `%s' must be a bytes but is %s" % (value, type(value))) 
184          self._dict[key] = value 
185          self._needs_sync = 0o001 
186   
187 -    def keys(self): 
188          return list(self._dict.keys()) 
189   
190 -    def has_key(self, key): 
191          return key in self._dict 
192   
193 -    def __contains__(self, key): 
194          return key in self._dict 
195   
196 -    def iterkeys(self): 
197          # Wrapping name in () prevents fixer from "fixing" this 
198          return (self._dict.iterkeys)() 
199   
200      __iter__ = iterkeys 
201   
202 -    def __len__(self): 
203          return len(self._dict) 
204   
205   
206 -def open(file, flag=None, mode=0o666): 
207      return dblite(file, flag, mode) 
208   
209   
210 -def _exercise(): 
211      db = open("tmp", "n") 
212      assert len(db) == 0 
213      db["foo"] = "bar" 
214      assert db["foo"] == "bar" 
215      db[unicode("ufoo")] = unicode("ubar") 
216      assert db[unicode("ufoo")] == unicode("ubar") 
217      db.sync() 
218      db = open("tmp", "c") 
219      assert len(db) == 2, len(db) 
220      assert db["foo"] == "bar" 
221      db["bar"] = "foo" 
222      assert db["bar"] == "foo" 
223      db[unicode("ubar")] = unicode("ufoo") 
224      assert db[unicode("ubar")] == unicode("ufoo") 
225      db.sync() 
226      db = open("tmp", "r") 
227      assert len(db) == 4, len(db) 
228      assert db["foo"] == "bar" 
229      assert db["bar"] == "foo" 
230      assert db[unicode("ufoo")] == unicode("ubar") 
231      assert db[unicode("ubar")] == unicode("ufoo") 
232      try: 
233          db.sync() 
234      except IOError as e: 
235          assert str(e) == "Read-only database: tmp.dblite" 
236      else: 
237          raise RuntimeError("IOError expected.") 
238      db = open("tmp", "w") 
239      assert len(db) == 4 
240      db["ping"] = "pong" 
241      db.sync() 
242      try: 
243          db[(1, 2)] = "tuple" 
244      except TypeError as e: 
245          assert str(e) == "key `(1, 2)' must be a string but is <type 'tuple'>", str(e) 
246      else: 
247          raise RuntimeError("TypeError exception expected") 
248      try: 
249          db["list"] = [1, 2] 
250      except TypeError as e: 
251          assert str(e) == "value `[1, 2]' must be a string but is <type 'list'>", str(e) 
252      else: 
253          raise RuntimeError("TypeError exception expected") 
254      db = open("tmp", "r") 
255      assert len(db) == 5 
256      db = open("tmp", "n") 
257      assert len(db) == 0 
258      dblite._open("tmp.dblite", "w") 
259      db = open("tmp", "r") 
260      dblite._open("tmp.dblite", "w").write("x") 
261      try: 
262          db = open("tmp", "r") 
263      except pickle.UnpicklingError: 
264          pass 
265      else: 
266          raise RuntimeError("pickle exception expected.") 
267      global ignore_corrupt_dbfiles 
268      ignore_corrupt_dbfiles = 2 
269      db = open("tmp", "r") 
270      assert len(db) == 0 
271      os.unlink("tmp.dblite") 
272      try: 
273          db = open("tmp", "w") 
274      except IOError as e: 
275          assert str(e) == "[Errno 2] No such file or directory: 'tmp.dblite'", str(e) 
276      else: 
277          raise RuntimeError("IOError expected.") 
278   
279   
280  if (__name__ == "__main__"): 
281      _exercise() 
282   
283  # Local Variables: 
284  # tab-width:4 
285  # indent-tabs-mode:nil 
286  # End: 
287  # vim: set expandtab tabstop=4 shiftwidth=4: 
288