OLD | NEW |
(Empty) | |
| 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(val
ue))) |
| 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: |
OLD | NEW |