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