]> arthur.barton.de Git - bup.git/blob - lib/bup/index.py
Add docstrings to lib/bup/helpers.py
[bup.git] / lib / bup / index.py
1 import os, stat, time, struct, tempfile
2 from bup.helpers import *
3
4 EMPTY_SHA = '\0'*20
5 FAKE_SHA = '\x01'*20
6 INDEX_HDR = 'BUPI\0\0\0\2'
7 INDEX_SIG = '!IIIIIQII20sHII'
8 ENTLEN = struct.calcsize(INDEX_SIG)
9 FOOTER_SIG = '!Q'
10 FOOTLEN = struct.calcsize(FOOTER_SIG)
11
12 IX_EXISTS = 0x8000        # file exists on filesystem
13 IX_HASHVALID = 0x4000     # the stored sha1 matches the filesystem
14 IX_SHAMISSING = 0x2000    # the stored sha1 object doesn't seem to exist
15
16 class Error(Exception):
17     pass
18
19
20 class Level:
21     def __init__(self, ename, parent):
22         self.parent = parent
23         self.ename = ename
24         self.list = []
25         self.count = 0
26
27     def write(self, f):
28         (ofs,n) = (f.tell(), len(self.list))
29         if self.list:
30             count = len(self.list)
31             #log('popping %r with %d entries\n' 
32             #    % (''.join(self.ename), count))
33             for e in self.list:
34                 e.write(f)
35             if self.parent:
36                 self.parent.count += count + self.count
37         return (ofs,n)
38
39
40 def _golevel(level, f, ename, newentry):
41     # close nodes back up the tree
42     assert(level)
43     while ename[:len(level.ename)] != level.ename:
44         n = BlankNewEntry(level.ename[-1])
45         n.flags |= IX_EXISTS
46         (n.children_ofs,n.children_n) = level.write(f)
47         level.parent.list.append(n)
48         level = level.parent
49
50     # create nodes down the tree
51     while len(level.ename) < len(ename):
52         level = Level(ename[:len(level.ename)+1], level)
53
54     # are we in precisely the right place?
55     assert(ename == level.ename)
56     n = newentry or BlankNewEntry(ename and level.ename[-1] or None)
57     (n.children_ofs,n.children_n) = level.write(f)
58     if level.parent:
59         level.parent.list.append(n)
60     level = level.parent
61
62     return level
63
64
65 class Entry:
66     def __init__(self, basename, name):
67         self.basename = str(basename)
68         self.name = str(name)
69         self.children_ofs = 0
70         self.children_n = 0
71
72     def __repr__(self):
73         return ("(%s,0x%04x,%d,%d,%d,%d,%d,%s/%s,0x%04x,0x%08x/%d)" 
74                 % (self.name, self.dev,
75                    self.ctime, self.mtime, self.uid, self.gid,
76                    self.size, oct(self.mode), oct(self.gitmode),
77                    self.flags, self.children_ofs, self.children_n))
78
79     def packed(self):
80         return struct.pack(INDEX_SIG,
81                            self.dev, self.ctime, self.mtime, 
82                            self.uid, self.gid, self.size, self.mode,
83                            self.gitmode, self.sha, self.flags,
84                            self.children_ofs, self.children_n)
85
86     def from_stat(self, st, tstart):
87         old = (self.dev, self.ctime, self.mtime,
88                self.uid, self.gid, self.size, self.flags & IX_EXISTS)
89         new = (st.st_dev, int(st.st_ctime), int(st.st_mtime),
90                st.st_uid, st.st_gid, st.st_size, IX_EXISTS)
91         self.dev = st.st_dev
92         self.ctime = int(st.st_ctime)
93         self.mtime = int(st.st_mtime)
94         self.uid = st.st_uid
95         self.gid = st.st_gid
96         self.size = st.st_size
97         self.mode = st.st_mode
98         self.flags |= IX_EXISTS
99         if int(st.st_ctime) >= tstart or old != new \
100               or self.sha == EMPTY_SHA or not self.gitmode:
101             self.invalidate()
102
103     def is_valid(self):
104         f = IX_HASHVALID|IX_EXISTS
105         return (self.flags & f) == f
106
107     def invalidate(self):
108         self.flags &= ~IX_HASHVALID
109
110     def validate(self, gitmode, sha):
111         assert(sha)
112         assert(gitmode)
113         self.gitmode = gitmode
114         self.sha = sha
115         self.flags |= IX_HASHVALID|IX_EXISTS
116
117     def exists(self):
118         return not self.is_deleted()
119
120     def sha_missing(self):
121         return (self.flags & IX_SHAMISSING) or not (self.flags & IX_HASHVALID)
122
123     def is_deleted(self):
124         return (self.flags & IX_EXISTS) == 0
125
126     def set_deleted(self):
127         if self.flags & IX_EXISTS:
128             self.flags &= ~(IX_EXISTS | IX_HASHVALID)
129
130     def is_real(self):
131         return not self.is_fake()
132
133     def is_fake(self):
134         return not self.ctime
135
136     def __cmp__(a, b):
137         return (cmp(a.name, b.name)
138                 or -cmp(a.is_valid(), b.is_valid())
139                 or -cmp(a.is_fake(), b.is_fake()))
140
141     def write(self, f):
142         f.write(self.basename + '\0' + self.packed())
143
144
145 class NewEntry(Entry):
146     def __init__(self, basename, name, dev, ctime, mtime, uid, gid,
147                  size, mode, gitmode, sha, flags, children_ofs, children_n):
148         Entry.__init__(self, basename, name)
149         (self.dev, self.ctime, self.mtime, self.uid, self.gid,
150          self.size, self.mode, self.gitmode, self.sha,
151          self.flags, self.children_ofs, self.children_n
152          ) = (dev, int(ctime), int(mtime), uid, gid,
153               size, mode, gitmode, sha, flags, children_ofs, children_n)
154
155
156 class BlankNewEntry(NewEntry):
157     def __init__(self, basename):
158         NewEntry.__init__(self, basename, basename,
159                           0, 0, 0, 0, 0, 0, 0,
160                           0, EMPTY_SHA, 0, 0, 0)
161
162
163 class ExistingEntry(Entry):
164     def __init__(self, parent, basename, name, m, ofs):
165         Entry.__init__(self, basename, name)
166         self.parent = parent
167         self._m = m
168         self._ofs = ofs
169         (self.dev, self.ctime, self.mtime, self.uid, self.gid,
170          self.size, self.mode, self.gitmode, self.sha,
171          self.flags, self.children_ofs, self.children_n
172          ) = struct.unpack(INDEX_SIG, str(buffer(m, ofs, ENTLEN)))
173
174     # effectively, we don't bother messing with IX_SHAMISSING if
175     # not IX_HASHVALID, since it's redundant, and repacking is more
176     # expensive than not repacking.
177     # This is implemented by having sha_missing() check IX_HASHVALID too.
178     def set_sha_missing(self, val):
179         val = val and 1 or 0
180         oldval = self.sha_missing() and 1 or 0
181         if val != oldval:
182             flag = val and IX_SHAMISSING or 0
183             newflags = (self.flags & (~IX_SHAMISSING)) | flag
184             self.flags = newflags
185             self.repack()
186
187     def unset_sha_missing(self, flag):
188         if self.flags & IX_SHAMISSING:
189             self.flags &= ~IX_SHAMISSING
190             self.repack()
191
192     def repack(self):
193         self._m[self._ofs:self._ofs+ENTLEN] = self.packed()
194         if self.parent and not self.is_valid():
195             self.parent.invalidate()
196             self.parent.repack()
197
198     def iter(self, name=None, wantrecurse=None):
199         dname = name
200         if dname and not dname.endswith('/'):
201             dname += '/'
202         ofs = self.children_ofs
203         assert(ofs <= len(self._m))
204         assert(self.children_n < 1000000)
205         for i in xrange(self.children_n):
206             eon = self._m.find('\0', ofs)
207             assert(eon >= 0)
208             assert(eon >= ofs)
209             assert(eon > ofs)
210             basename = str(buffer(self._m, ofs, eon-ofs))
211             child = ExistingEntry(self, basename, self.name + basename,
212                                   self._m, eon+1)
213             if (not dname
214                  or child.name.startswith(dname)
215                  or child.name.endswith('/') and dname.startswith(child.name)):
216                 if not wantrecurse or wantrecurse(child):
217                     for e in child.iter(name=name, wantrecurse=wantrecurse):
218                         yield e
219             if not name or child.name == name or child.name.startswith(dname):
220                 yield child
221             ofs = eon + 1 + ENTLEN
222
223     def __iter__(self):
224         return self.iter()
225             
226
227 class Reader:
228     def __init__(self, filename):
229         self.filename = filename
230         self.m = ''
231         self.writable = False
232         self.count = 0
233         f = None
234         try:
235             f = open(filename, 'r+')
236         except IOError, e:
237             if e.errno == errno.ENOENT:
238                 pass
239             else:
240                 raise
241         if f:
242             b = f.read(len(INDEX_HDR))
243             if b != INDEX_HDR:
244                 log('warning: %s: header: expected %r, got %r'
245                                  % (filename, INDEX_HDR, b))
246             else:
247                 st = os.fstat(f.fileno())
248                 if st.st_size:
249                     self.m = mmap_readwrite(f)
250                     self.writable = True
251                     self.count = struct.unpack(FOOTER_SIG,
252                           str(buffer(self.m, st.st_size-FOOTLEN, FOOTLEN)))[0]
253
254     def __del__(self):
255         self.close()
256
257     def __len__(self):
258         return int(self.count)
259
260     def forward_iter(self):
261         ofs = len(INDEX_HDR)
262         while ofs+ENTLEN <= len(self.m)-FOOTLEN:
263             eon = self.m.find('\0', ofs)
264             assert(eon >= 0)
265             assert(eon >= ofs)
266             assert(eon > ofs)
267             basename = str(buffer(self.m, ofs, eon-ofs))
268             yield ExistingEntry(None, basename, basename, self.m, eon+1)
269             ofs = eon + 1 + ENTLEN
270
271     def iter(self, name=None, wantrecurse=None):
272         if len(self.m) > len(INDEX_HDR)+ENTLEN:
273             dname = name
274             if dname and not dname.endswith('/'):
275                 dname += '/'
276             root = ExistingEntry(None, '/', '/',
277                                  self.m, len(self.m)-FOOTLEN-ENTLEN)
278             for sub in root.iter(name=name, wantrecurse=wantrecurse):
279                 yield sub
280             if not dname or dname == root.name:
281                 yield root
282
283     def __iter__(self):
284         return self.iter()
285
286     def exists(self):
287         return self.m
288
289     def save(self):
290         if self.writable and self.m:
291             self.m.flush()
292
293     def close(self):
294         self.save()
295         if self.writable and self.m:
296             self.m.close()
297             self.m = None
298             self.writable = False
299
300     def filter(self, prefixes, wantrecurse=None):
301         for (rp, path) in reduce_paths(prefixes):
302             for e in self.iter(rp, wantrecurse=wantrecurse):
303                 assert(e.name.startswith(rp))
304                 name = path + e.name[len(rp):]
305                 yield (name, e)
306
307
308 # FIXME: this function isn't very generic, because it splits the filename
309 # in an odd way and depends on a terminating '/' to indicate directories.
310 def pathsplit(p):
311     """Split a path into a list of elements of the file system hierarchy."""
312     l = p.split('/')
313     l = [i+'/' for i in l[:-1]] + l[-1:]
314     if l[-1] == '':
315         l.pop()  # extra blank caused by terminating '/'
316     return l
317
318
319 class Writer:
320     def __init__(self, filename):
321         self.rootlevel = self.level = Level([], None)
322         self.f = None
323         self.count = 0
324         self.lastfile = None
325         self.filename = None
326         self.filename = filename = realpath(filename)
327         (dir,name) = os.path.split(filename)
328         (ffd,self.tmpname) = tempfile.mkstemp('.tmp', filename, dir)
329         self.f = os.fdopen(ffd, 'wb', 65536)
330         self.f.write(INDEX_HDR)
331
332     def __del__(self):
333         self.abort()
334
335     def abort(self):
336         f = self.f
337         self.f = None
338         if f:
339             f.close()
340             os.unlink(self.tmpname)
341
342     def flush(self):
343         if self.level:
344             self.level = _golevel(self.level, self.f, [], None)
345             self.count = self.rootlevel.count
346             if self.count:
347                 self.count += 1
348             self.f.write(struct.pack(FOOTER_SIG, self.count))
349             self.f.flush()
350         assert(self.level == None)
351
352     def close(self):
353         self.flush()
354         f = self.f
355         self.f = None
356         if f:
357             f.close()
358             os.rename(self.tmpname, self.filename)
359
360     def _add(self, ename, entry):
361         if self.lastfile and self.lastfile <= ename:
362             raise Error('%r must come before %r' 
363                              % (''.join(e.name), ''.join(self.lastfile)))
364             self.lastfile = e.name
365         self.level = _golevel(self.level, self.f, ename, entry)
366
367     def add(self, name, st, hashgen = None):
368         endswith = name.endswith('/')
369         ename = pathsplit(name)
370         basename = ename[-1]
371         #log('add: %r %r\n' % (basename, name))
372         flags = IX_EXISTS
373         sha = None
374         if hashgen:
375             (gitmode, sha) = hashgen(name)
376             flags |= IX_HASHVALID
377         else:
378             (gitmode, sha) = (0, EMPTY_SHA)
379         if st:
380             isdir = stat.S_ISDIR(st.st_mode)
381             assert(isdir == endswith)
382             e = NewEntry(basename, name, st.st_dev, int(st.st_ctime),
383                          int(st.st_mtime), st.st_uid, st.st_gid,
384                          st.st_size, st.st_mode, gitmode, sha, flags,
385                          0, 0)
386         else:
387             assert(endswith)
388             e = BlankNewEntry(basename)
389             e.gitmode = gitmode
390             e.sha = sha
391             e.flags = flags
392         self._add(ename, e)
393
394     def add_ixentry(self, e):
395         e.children_ofs = e.children_n = 0
396         self._add(pathsplit(e.name), e)
397
398     def new_reader(self):
399         self.flush()
400         return Reader(self.tmpname)
401
402
403 def reduce_paths(paths):
404     xpaths = []
405     for p in paths:
406         rp = realpath(p)
407         st = os.lstat(rp)
408         if stat.S_ISDIR(st.st_mode):
409             rp = slashappend(rp)
410             p = slashappend(p)
411         xpaths.append((rp, p))
412     xpaths.sort()
413
414     paths = []
415     prev = None
416     for (rp, p) in xpaths:
417         if prev and (prev == rp 
418                      or (prev.endswith('/') and rp.startswith(prev))):
419             continue # already superceded by previous path
420         paths.append((rp, p))
421         prev = rp
422     paths.sort(reverse=True)
423     return paths
424
425
426 class MergeIter:
427     def __init__(self, iters):
428         self.iters = iters
429
430     def __len__(self):
431         # FIXME: doesn't remove duplicated entries between iters.
432         # That only happens for parent directories, but will mean the
433         # actual iteration returns fewer entries than this function counts.
434         return sum(len(it) for it in self.iters)
435
436     def __iter__(self):
437         total = len(self)
438         l = [iter(it) for it in self.iters]
439         l = [(next(it),it) for it in l]
440         l = filter(lambda x: x[0], l)
441         count = 0
442         lastname = None
443         while l:
444             if not (count % 1024):
445                 progress('bup: merging indexes (%d/%d)\r' % (count, total))
446             l.sort()
447             (e,it) = l.pop()
448             if not e:
449                 continue
450             if e.name != lastname:
451                 yield e
452                 lastname = e.name
453             n = next(it)
454             if n:
455                 l.append((n,it))
456             count += 1
457         log('bup: merging indexes (%d/%d), done.\n' % (count, total))