index.py

   1 import os, stat, time, struct, tempfile, mmap
   2 from helpers import *
   3
   4 EMPTY_SHA = '\0'*20
   5 FAKE_SHA = '\x01'*20
   6 INDEX_HDR = 'BUPI\0\0\0\1'
   7 INDEX_SIG = '!IIIIIQII20sH'
   8 ENTLEN = struct.calcsize(INDEX_SIG)
   9
  10 IX_EXISTS = 0x8000
  11 IX_HASHVALID = 0x4000
  12
  13 class Error(Exception):
  14     pass
  15
  16
  17 class Entry:
  18     def __init__(self, name, m, ofs, tstart):
  19         self._m = m
  20         self._ofs = ofs
  21         self.name = str(name)
  22         self.tstart = tstart
  23         (self.dev, self.ctime, self.mtime, self.uid, self.gid,
  24          self.size, self.mode, self.gitmode, self.sha,
  25          self.flags) = struct.unpack(INDEX_SIG, str(buffer(m, ofs, ENTLEN)))
  26
  27     def __repr__(self):
  28         return ("(%s,0x%04x,%d,%d,%d,%d,%d,0x%04x)"
  29                 % (self.name, self.dev,
  30                    self.ctime, self.mtime, self.uid, self.gid,
  31                    self.size, self.flags))
  32
  33     def packed(self):
  34         return struct.pack(INDEX_SIG,
  35                            self.dev, self.ctime, self.mtime,
  36                            self.uid, self.gid, self.size, self.mode,
  37                            self.gitmode, self.sha, self.flags)
  38
  39     def repack(self):
  40         self._m[self._ofs:self._ofs+ENTLEN] = self.packed()
  41
  42     def from_stat(self, st):
  43         old = (self.dev, self.ctime, self.mtime,
  44                self.uid, self.gid, self.size, self.flags & IX_EXISTS)
  45         new = (st.st_dev, int(st.st_ctime), int(st.st_mtime),
  46                st.st_uid, st.st_gid, st.st_size, IX_EXISTS)
  47         self.dev = st.st_dev
  48         self.ctime = int(st.st_ctime)
  49         self.mtime = int(st.st_mtime)
  50         self.uid = st.st_uid
  51         self.gid = st.st_gid
  52         self.size = st.st_size
  53         self.mode = st.st_mode
  54         self.flags |= IX_EXISTS
  55         if int(st.st_ctime) >= self.tstart or old != new:
  56             self.flags &= ~IX_HASHVALID
  57             return 1  # dirty
  58         else:
  59             return 0  # not dirty
  60
  61     def validate(self, sha):
  62         assert(sha)
  63         self.sha = sha
  64         self.flags |= IX_HASHVALID
  65
  66     def __cmp__(a, b):
  67         return cmp(a.name, b.name)
  68
  69
  70 class Reader:
  71     def __init__(self, filename):
  72         self.filename = filename
  73         self.m = ''
  74         self.writable = False
  75         f = None
  76         try:
  77             f = open(filename, 'r+')
  78         except IOError, e:
  79             if e.errno == errno.ENOENT:
  80                 pass
  81             else:
  82                 raise
  83         if f:
  84             b = f.read(len(INDEX_HDR))
  85             if b != INDEX_HDR:
  86                 raise Error('%s: header: expected %r, got %r'
  87                                  % (filename, INDEX_HDR, b))
  88             st = os.fstat(f.fileno())
  89             if st.st_size:
  90                 self.m = mmap.mmap(f.fileno(), 0,
  91                                    mmap.MAP_SHARED,
  92                                    mmap.PROT_READ|mmap.PROT_WRITE)
  93                 f.close()  # map will persist beyond file close
  94                 self.writable = True
  95
  96     def __del__(self):
  97         self.close()
  98
  99     def __iter__(self):
 100         tstart = int(time.time())
 101         ofs = len(INDEX_HDR)
 102         while ofs < len(self.m):
 103             eon = self.m.find('\0', ofs)
 104             assert(eon >= 0)
 105             yield Entry(buffer(self.m, ofs, eon-ofs),
 106                           self.m, eon+1, tstart = tstart)
 107             ofs = eon + 1 + ENTLEN
 108
 109     def save(self):
 110         if self.writable:
 111             self.m.flush()
 112
 113     def close(self):
 114         self.save()
 115         if self.writable:
 116             self.m.close()
 117             self.writable = False
 118
 119     def filter(self, prefixes):
 120         #log("filtering %r\n" % prefixes)
 121         paths = reduce_paths(prefixes)
 122         #log("filtering %r\n" % paths)
 123         pi = iter(paths)
 124         (rpin, pin) = pi.next()
 125         for ent in self:
 126             #log('checking %r vs %r\n' % (ent.name, rpin))
 127             while ent.name < rpin:
 128                 try:
 129                     (rpin, pin) = pi.next()
 130                 except StopIteration:
 131                     return  # no more files can possibly match
 132             if not ent.name.startswith(rpin):
 133                 continue   # not interested
 134             else:
 135                 name = pin + ent.name[len(rpin):]
 136                 yield (name, ent)
 137
 138
 139 # Read all the iters in order; when more than one iter has the same entry,
 140 # the *later* iter in the list wins.  (ie. more recent iter entries replace
 141 # older ones)
 142 def _last_writer_wins_iter(iters):
 143     l = []
 144     for e in iters:
 145         it = iter(e)
 146         try:
 147             l.append([it.next(), it])
 148         except StopIteration:
 149             pass
 150     del iters  # to avoid accidents
 151     while l:
 152         l.sort()
 153         mv = l[0][0]
 154         mi = []
 155         for (i,(v,it)) in enumerate(l):
 156             #log('(%d) considering %d: %r\n' % (len(l), i, v))
 157             if v > mv:
 158                 mv = v
 159                 mi = [i]
 160             elif v == mv:
 161                 mi.append(i)
 162         yield mv
 163         for i in mi:
 164             try:
 165                 l[i][0] = l[i][1].next()
 166             except StopIteration:
 167                 l[i] = None
 168         l = filter(None, l)
 169
 170
 171 class Writer:
 172     def __init__(self, filename):
 173         self.f = None
 174         self.count = 0
 175         self.lastfile = None
 176         self.filename = None
 177         self.filename = filename = realpath(filename)
 178         (dir,name) = os.path.split(filename)
 179         (ffd,self.tmpname) = tempfile.mkstemp('.tmp', filename, dir)
 180         self.f = os.fdopen(ffd, 'wb', 65536)
 181         self.f.write(INDEX_HDR)
 182
 183     def __del__(self):
 184         self.abort()
 185
 186     def abort(self):
 187         f = self.f
 188         self.f = None
 189         if f:
 190             f.close()
 191             os.unlink(self.tmpname)
 192
 193     def close(self):
 194         f = self.f
 195         self.f = None
 196         if f:
 197             f.close()
 198             try:
 199                 os.rename(self.tmpname, self.filename)
 200             except OSError:
 201                 if os.path.exists(self.filename):
 202                     os.unlink(self.filename)
 203                 os.rename(self.tmpname, self.filename)
 204
 205     def _write(self, data):
 206         self.f.write(data)
 207         self.count += 1
 208
 209     def add(self, name, st, hashgen=None):
 210         #log('ADDING %r\n' % name)
 211         if self.lastfile:
 212             assert(cmp(self.lastfile, name) > 0) # reverse order only
 213         self.lastfile = name
 214         flags = IX_EXISTS
 215         sha = None
 216         if hashgen:
 217             (gitmode, sha) = hashgen(name)
 218             if sha:
 219                 flags |= IX_HASHVALID
 220         else:
 221             (gitmode, sha) = (0, EMPTY_SHA)
 222         data = name + '\0' + \
 223             struct.pack(INDEX_SIG, st.st_dev, int(st.st_ctime),
 224                         int(st.st_mtime), st.st_uid, st.st_gid,
 225                         st.st_size, st.st_mode, gitmode, sha, flags)
 226         self._write(data)
 227
 228     def add_ixentry(self, e):
 229         if self.lastfile and self.lastfile <= e.name:
 230             raise Error('%r must come before %r'
 231                              % (e.name, self.lastfile))
 232         self.lastfile = e.name
 233         data = e.name + '\0' + e.packed()
 234         self._write(data)
 235
 236     def new_reader(self):
 237         self.f.flush()
 238         return Reader(self.tmpname)
 239
 240
 241 # like os.path.realpath, but doesn't follow a symlink for the last element.
 242 # (ie. if 'p' itself is itself a symlink, this one won't follow it)
 243 def realpath(p):
 244     try:
 245         st = os.lstat(p)
 246     except OSError:
 247         st = None
 248     if st and stat.S_ISLNK(st.st_mode):
 249         (dir, name) = os.path.split(p)
 250         dir = os.path.realpath(dir)
 251         out = os.path.join(dir, name)
 252     else:
 253         out = os.path.realpath(p)
 254     #log('realpathing:%r,%r\n' % (p, out))
 255     return out
 256
 257
 258 def reduce_paths(paths):
 259     xpaths = []
 260     for p in paths:
 261         rp = realpath(p)
 262         st = os.lstat(rp)
 263         if stat.S_ISDIR(st.st_mode):
 264             rp = slashappend(rp)
 265             p = slashappend(p)
 266         xpaths.append((rp, p))
 267     xpaths.sort()
 268
 269     paths = []
 270     prev = None
 271     for (rp, p) in xpaths:
 272         if prev and (prev == rp
 273                      or (prev.endswith('/') and rp.startswith(prev))):
 274             continue # already superceded by previous path
 275         paths.append((rp, p))
 276         prev = rp
 277     paths.sort(reverse=True)
 278     return paths
 279