1
2
3
4 import cPickle
5 import time
6 import shutil
7 import os
8 import types
9 import __builtin__
10
11 keep_all_files = 00000
12 ignore_corrupt_dbfiles = 0
13
15 print "Warning: Discarding corrupt database:", filename
16
17 if hasattr(types, 'UnicodeType'):
19 t = type(s)
20 return t is types.StringType or t is types.UnicodeType
21 else:
23 return type(s) is types.StringType
24
25 try:
26 unicode('a')
27 except NameError:
29
30 dblite_suffix = '.dblite'
31 tmp_suffix = '.tmp'
32
34
35
36
37
38
39
40
41
42
43
44
45 _open = __builtin__.open
46 _cPickle_dump = cPickle.dump
47 _os_chmod = os.chmod
48 try:
49 _os_chown = os.chown
50 except AttributeError:
51 _os_chown = None
52 _os_rename = os.rename
53 _os_unlink = os.unlink
54 _shutil_copyfile = shutil.copyfile
55 _time_time = time.time
56
57 - def __init__(self, file_base_name, flag, mode):
58 assert flag in (None, "r", "w", "c", "n")
59 if (flag is None): flag = "r"
60 base, ext = os.path.splitext(file_base_name)
61 if ext == dblite_suffix:
62
63 self._file_name = file_base_name
64 self._tmp_name = base + tmp_suffix
65 else:
66 self._file_name = file_base_name + dblite_suffix
67 self._tmp_name = file_base_name + tmp_suffix
68 self._flag = flag
69 self._mode = mode
70 self._dict = {}
71 self._needs_sync = 00000
72 if self._os_chown is not None and (os.geteuid()==0 or os.getuid()==0):
73
74 try:
75 statinfo = os.stat(self._file_name)
76 self._chown_to = statinfo.st_uid
77 self._chgrp_to = statinfo.st_gid
78 except OSError, e:
79
80
81 self._chown_to = int(os.environ.get('SUDO_UID', -1))
82 self._chgrp_to = int(os.environ.get('SUDO_GID', -1))
83 else:
84 self._chown_to = -1
85 self._chgrp_to = -1
86 if (self._flag == "n"):
87 self._open(self._file_name, "wb", self._mode)
88 else:
89 try:
90 f = self._open(self._file_name, "rb")
91 except IOError, e:
92 if (self._flag != "c"):
93 raise e
94 self._open(self._file_name, "wb", self._mode)
95 else:
96 p = f.read()
97 if (len(p) > 0):
98 try:
99 self._dict = cPickle.loads(p)
100 except (cPickle.UnpicklingError, EOFError):
101 if (ignore_corrupt_dbfiles == 0): raise
102 if (ignore_corrupt_dbfiles == 1):
103 corruption_warning(self._file_name)
104
106 if (self._needs_sync):
107 self.sync()
108
110 self._check_writable()
111 f = self._open(self._tmp_name, "wb", self._mode)
112 self._cPickle_dump(self._dict, f, 1)
113 f.close()
114
115
116
117
118
119
120 try: self._os_chmod(self._file_name, 0777)
121 except OSError: pass
122 self._os_unlink(self._file_name)
123 self._os_rename(self._tmp_name, self._file_name)
124 if self._os_chown is not None and self._chown_to > 0:
125 try:
126 self._os_chown(self._file_name, self._chown_to, self._chgrp_to)
127 except OSError:
128 pass
129 self._needs_sync = 00000
130 if (keep_all_files):
131 self._shutil_copyfile(
132 self._file_name,
133 self._file_name + "_" + str(int(self._time_time())))
134
136 if (self._flag == "r"):
137 raise IOError("Read-only database: %s" % self._file_name)
138
140 return self._dict[key]
141
143 self._check_writable()
144 if (not is_string(key)):
145 raise TypeError, "key `%s' must be a string but is %s" % (key, type(key))
146 if (not is_string(value)):
147 raise TypeError, "value `%s' must be a string but is %s" % (value, type(value))
148 self._dict[key] = value
149 self._needs_sync = 0001
150
152 return self._dict.keys()
153
155 return key in self._dict
156
158 return key in self._dict
159
162
163 __iter__ = iterkeys
164
166 return len(self._dict)
167
168 -def open(file, flag=None, mode=0666):
169 return dblite(file, flag, mode)
170
172 db = open("tmp", "n")
173 assert len(db) == 0
174 db["foo"] = "bar"
175 assert db["foo"] == "bar"
176 db[unicode("ufoo")] = unicode("ubar")
177 assert db[unicode("ufoo")] == unicode("ubar")
178 db.sync()
179 db = open("tmp", "c")
180 assert len(db) == 2, len(db)
181 assert db["foo"] == "bar"
182 db["bar"] = "foo"
183 assert db["bar"] == "foo"
184 db[unicode("ubar")] = unicode("ufoo")
185 assert db[unicode("ubar")] == unicode("ufoo")
186 db.sync()
187 db = open("tmp", "r")
188 assert len(db) == 4, len(db)
189 assert db["foo"] == "bar"
190 assert db["bar"] == "foo"
191 assert db[unicode("ufoo")] == unicode("ubar")
192 assert db[unicode("ubar")] == unicode("ufoo")
193 try:
194 db.sync()
195 except IOError, e:
196 assert str(e) == "Read-only database: tmp.dblite"
197 else:
198 raise RuntimeError, "IOError expected."
199 db = open("tmp", "w")
200 assert len(db) == 4
201 db["ping"] = "pong"
202 db.sync()
203 try:
204 db[(1,2)] = "tuple"
205 except TypeError, e:
206 assert str(e) == "key `(1, 2)' must be a string but is <type 'tuple'>", str(e)
207 else:
208 raise RuntimeError, "TypeError exception expected"
209 try:
210 db["list"] = [1,2]
211 except TypeError, e:
212 assert str(e) == "value `[1, 2]' must be a string but is <type 'list'>", str(e)
213 else:
214 raise RuntimeError, "TypeError exception expected"
215 db = open("tmp", "r")
216 assert len(db) == 5
217 db = open("tmp", "n")
218 assert len(db) == 0
219 _open("tmp.dblite", "w")
220 db = open("tmp", "r")
221 _open("tmp.dblite", "w").write("x")
222 try:
223 db = open("tmp", "r")
224 except cPickle.UnpicklingError:
225 pass
226 else:
227 raise RuntimeError, "cPickle exception expected."
228 global ignore_corrupt_dbfiles
229 ignore_corrupt_dbfiles = 2
230 db = open("tmp", "r")
231 assert len(db) == 0
232 os.unlink("tmp.dblite")
233 try:
234 db = open("tmp", "w")
235 except IOError, e:
236 assert str(e) == "[Errno 2] No such file or directory: 'tmp.dblite'", str(e)
237 else:
238 raise RuntimeError, "IOError expected."
239 print "OK"
240
241 if (__name__ == "__main__"):
242 _exercise()
243
244
245
246
247
248
249