1# MIT License 2# 3# Copyright The SCons Foundation 4# 5# Permission is hereby granted, free of charge, to any person obtaining 6# a copy of this software and associated documentation files (the 7# "Software"), to deal in the Software without restriction, including 8# without limitation the rights to use, copy, modify, merge, publish, 9# distribute, sublicense, and/or sell copies of the Software, and to 10# permit persons to whom the Software is furnished to do so, subject to 11# the following conditions: 12# 13# The above copyright notice and this permission notice shall be included 14# in all copies or substantial portions of the Software. 15# 16# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY 17# KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE 18# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 19# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 20# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 21# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 22# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 23 24""" 25dblite.py module contributed by Ralf W. Grosse-Kunstleve. 26Extended for Unicode by Steven Knight. 27""" 28 29import os 30import pickle 31import shutil 32import time 33 34from SCons.compat import PICKLE_PROTOCOL 35 36KEEP_ALL_FILES = False 37IGNORE_CORRUPT_DBFILES = False 38 39 40def corruption_warning(filename): 41 """Local warning for corrupt db. 42 43 Used for self-tests. SCons overwrites this with a 44 different warning function in SConsign.py. 45 """ 46 print("Warning: Discarding corrupt database:", filename) 47 48DBLITE_SUFFIX = '.dblite' 49TMP_SUFFIX = '.tmp' 50 51 52class dblite: 53 """ 54 Squirrel away references to the functions in various modules 55 that we'll use when our __del__() method calls our sync() method 56 during shutdown. We might get destroyed when Python is in the midst 57 of tearing down the different modules we import in an essentially 58 arbitrary order, and some of the various modules's global attributes 59 may already be wiped out from under us. 60 61 See the discussion at: 62 http://mail.python.org/pipermail/python-bugs-list/2003-March/016877.html 63 64 """ 65 66 _open = open 67 _pickle_dump = staticmethod(pickle.dump) 68 _pickle_protocol = PICKLE_PROTOCOL 69 70 try: 71 _os_chown = os.chown 72 except AttributeError: 73 _os_chown = None 74 _os_replace = os.replace 75 _os_chmod = os.chmod 76 _shutil_copyfile = shutil.copyfile 77 _time_time = time.time 78 79 def __init__(self, file_base_name, flag, mode): 80 assert flag in (None, "r", "w", "c", "n") 81 if flag is None: 82 flag = "r" 83 84 base, ext = os.path.splitext(file_base_name) 85 if ext == DBLITE_SUFFIX: 86 # There's already a suffix on the file name, don't add one. 87 self._file_name = file_base_name 88 self._tmp_name = base + TMP_SUFFIX 89 else: 90 self._file_name = file_base_name + DBLITE_SUFFIX 91 self._tmp_name = file_base_name + TMP_SUFFIX 92 93 self._flag = flag 94 self._mode = mode 95 self._dict = {} 96 self._needs_sync = False 97 98 if self._os_chown is not None and (os.geteuid() == 0 or os.getuid() == 0): 99 # running as root; chown back to current owner/group when done 100 try: 101 statinfo = os.stat(self._file_name) 102 self._chown_to = statinfo.st_uid 103 self._chgrp_to = statinfo.st_gid 104 except OSError: 105 # db file doesn't exist yet. 106 # Check os.environ for SUDO_UID, use if set 107 self._chown_to = int(os.environ.get('SUDO_UID', -1)) 108 self._chgrp_to = int(os.environ.get('SUDO_GID', -1)) 109 else: 110 self._chown_to = -1 # don't chown 111 self._chgrp_to = -1 # don't chgrp 112 113 if self._flag == "n": 114 with self._open(self._file_name, "wb", self._mode): 115 pass # just make sure it exists 116 else: 117 try: 118 f = self._open(self._file_name, "rb") 119 except IOError as e: 120 if self._flag != "c": 121 raise e 122 with self._open(self._file_name, "wb", self._mode): 123 pass # just make sure it exists 124 else: 125 p = f.read() 126 f.close() 127 if len(p) > 0: 128 try: 129 self._dict = pickle.loads(p, encoding='bytes') 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: 135 corruption_warning(self._file_name) 136 else: 137 raise 138 139 def close(self): 140 if self._needs_sync: 141 self.sync() 142 143 def __del__(self): 144 self.close() 145 146 def sync(self): 147 self._check_writable() 148 with self._open(self._tmp_name, "wb", self._mode) as f: 149 self._pickle_dump(self._dict, f, self._pickle_protocol) 150 151 try: 152 self._os_replace(self._tmp_name, self._file_name) 153 except PermissionError: 154 # If we couldn't replace due to perms, try to change and retry. 155 # This is mainly for Windows - on POSIX the file permissions 156 # don't matter, the os.replace would have worked anyway. 157 # We're giving up if the retry fails, just let the Python 158 # exception abort us. 159 try: 160 self._os_chmod(self._file_name, 0o777) 161 except PermissionError: 162 pass 163 self._os_replace(self._tmp_name, self._file_name) 164 165 if self._os_chown is not None and self._chown_to > 0: # don't chown to root or -1 166 try: 167 self._os_chown(self._file_name, self._chown_to, self._chgrp_to) 168 except OSError: 169 pass 170 171 self._needs_sync = False 172 if KEEP_ALL_FILES: 173 self._shutil_copyfile( 174 self._file_name, 175 self._file_name + "_" + str(int(self._time_time())) 176 ) 177 178 def _check_writable(self): 179 if self._flag == "r": 180 raise IOError("Read-only database: %s" % self._file_name) 181 182 def __getitem__(self, key): 183 return self._dict[key] 184 185 def __setitem__(self, key, value): 186 self._check_writable() 187 188 if not isinstance(key, str): 189 raise TypeError("key `%s' must be a string but is %s" % (key, type(key))) 190 191 if not isinstance(value, bytes): 192 raise TypeError("value `%s' must be a bytes but is %s" % (value, type(value))) 193 194 self._dict[key] = value 195 self._needs_sync = True 196 197 def keys(self): 198 return list(self._dict.keys()) 199 200 def __contains__(self, key): 201 return key in self._dict 202 203 def __iter__(self): 204 return iter(self._dict) 205 206 def __len__(self): 207 return len(self._dict) 208 209 210def open(file, flag=None, mode=0o666): 211 return dblite(file, flag, mode) 212 213 214def _exercise(): 215 db = open("tmp", "n") 216 assert len(db) == 0 217 db["foo"] = b"bar" 218 assert db["foo"] == b"bar" 219 db.sync() 220 221 db = open("tmp", "c") 222 assert len(db) == 1, len(db) 223 assert db["foo"] == b"bar" 224 db["bar"] = b"foo" 225 assert db["bar"] == b"foo" 226 db.sync() 227 228 db = open("tmp", "r") 229 assert len(db) == 2, len(db) 230 assert db["foo"] == b"bar" 231 assert db["bar"] == b"foo" 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) == 2, len(db) 240 db["ping"] = b"pong" 241 db.sync() 242 243 try: 244 db[(1, 2)] = "tuple" 245 except TypeError as e: 246 assert str(e) == "key `(1, 2)' must be a string but is <class 'tuple'>", str(e) 247 else: 248 raise RuntimeError("TypeError exception expected") 249 250 try: 251 db["list"] = [1, 2] 252 except TypeError as e: 253 assert str(e) == "value `[1, 2]' must be a bytes but is <class 'list'>", str(e) 254 else: 255 raise RuntimeError("TypeError exception expected") 256 257 db = open("tmp", "r") 258 assert len(db) == 3, len(db) 259 260 db = open("tmp", "n") 261 assert len(db) == 0, len(db) 262 dblite._open("tmp.dblite", "w") 263 264 db = open("tmp", "r") 265 dblite._open("tmp.dblite", "w").write("x") 266 try: 267 db = open("tmp", "r") 268 except pickle.UnpicklingError: 269 pass 270 else: 271 raise RuntimeError("pickle exception expected.") 272 273 global IGNORE_CORRUPT_DBFILES 274 IGNORE_CORRUPT_DBFILES = True 275 db = open("tmp", "r") 276 assert len(db) == 0, len(db) 277 os.unlink("tmp.dblite") 278 try: 279 db = open("tmp", "w") 280 except IOError as e: 281 assert str(e) == "[Errno 2] No such file or directory: 'tmp.dblite'", str(e) 282 else: 283 raise RuntimeError("IOError expected.") 284 285 print("Completed _exercise()") 286 287 288if __name__ == "__main__": 289 _exercise() 290 291# Local Variables: 292# tab-width:4 293# indent-tabs-mode:nil 294# End: 295# vim: set expandtab tabstop=4 shiftwidth=4: 296