New dulwich.pack.MemoryPackIndex class.
[jelmer/dulwich-libgit2.git] / dulwich / pack.py
1 # pack.py -- For dealing with packed git objects.
2 # Copyright (C) 2007 James Westby <jw+debian@jameswestby.net>
3 # Copyright (C) 2008-2009 Jelmer Vernooij <jelmer@samba.org>
4 #
5 # This program is free software; you can redistribute it and/or
6 # modify it under the terms of the GNU General Public License
7 # as published by the Free Software Foundation; version 2
8 # of the License or (at your option) a later version.
9 #
10 # This program is distributed in the hope that it will be useful,
11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 # GNU General Public License for more details.
14 #
15 # You should have received a copy of the GNU General Public License
16 # along with this program; if not, write to the Free Software
17 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
18 # MA  02110-1301, USA.
19
20 """Classes for dealing with packed git objects.
21
22 A pack is a compact representation of a bunch of objects, stored
23 using deltas where possible.
24
25 They have two parts, the pack file, which stores the data, and an index
26 that tells you where the data is.
27
28 To find an object you look in all of the index files 'til you find a
29 match for the object name. You then use the pointer got from this as
30 a pointer in to the corresponding packfile.
31 """
32
33 try:
34     from collections import defaultdict
35 except ImportError:
36     from misc import defaultdict
37
38 from cStringIO import (
39     StringIO,
40     )
41 from collections import (
42     deque,
43     )
44 import difflib
45 from itertools import (
46     chain,
47     imap,
48     izip,
49     )
50 import mmap
51 import os
52 import struct
53 try:
54     from struct import unpack_from
55 except ImportError:
56     from dulwich.misc import unpack_from
57 import sys
58 import zlib
59
60 from dulwich.errors import (
61     ApplyDeltaError,
62     ChecksumMismatch,
63     )
64 from dulwich.file import GitFile
65 from dulwich.lru_cache import (
66     LRUSizeCache,
67     )
68 from dulwich.misc import (
69     make_sha,
70     SEEK_END,
71     )
72 from dulwich.objects import (
73     ShaFile,
74     hex_to_sha,
75     sha_to_hex,
76     object_header,
77     )
78
79 supports_mmap_offset = (sys.version_info[0] >= 3 or
80         (sys.version_info[0] == 2 and sys.version_info[1] >= 6))
81
82
83 OFS_DELTA = 6
84 REF_DELTA = 7
85
86 DELTA_TYPES = (OFS_DELTA, REF_DELTA)
87
88
89 def take_msb_bytes(read):
90     """Read bytes marked with most significant bit.
91
92     :param read: Read function
93     """
94     ret = []
95     while len(ret) == 0 or ret[-1] & 0x80:
96         ret.append(ord(read(1)))
97     return ret
98
99
100 def read_zlib_chunks(read_some, dec_size, buffer_size=4096):
101     """Read zlib data from a buffer.
102
103     This function requires that the buffer have additional data following the
104     compressed data, which is guaranteed to be the case for git pack files.
105
106     :param read_some: Read function that returns at least one byte, but may
107         return less than the requested size
108     :param dec_size: Expected size of the decompressed buffer
109     :param buffer_size: Size of the read buffer
110     :return: Tuple with list of chunks, length of compressed data length and
111         and unused read data.
112     :raise zlib.error: if a decompression error occurred.
113     """
114     if dec_size <= -1:
115         raise ValueError("non-negative zlib data stream size expected")
116     obj = zlib.decompressobj()
117     ret = []
118     fed = 0
119     size = 0
120     while obj.unused_data == "":
121         add = read_some(buffer_size)
122         if not add:
123             raise zlib.error("EOF before end of zlib stream")
124         fed += len(add)
125         decomp = obj.decompress(add)
126         size += len(decomp)
127         ret.append(decomp)
128     if size != dec_size:
129         raise zlib.error("decompressed data does not match expected size")
130     comp_len = fed - len(obj.unused_data)
131     return ret, comp_len, obj.unused_data
132
133
134 def iter_sha1(iter):
135     """Return the hexdigest of the SHA1 over a set of names.
136
137     :param iter: Iterator over string objects
138     :return: 40-byte hex sha1 digest
139     """
140     sha1 = make_sha()
141     for name in iter:
142         sha1.update(name)
143     return sha1.hexdigest()
144
145
146 def load_pack_index(path):
147     """Load an index file by path.
148
149     :param filename: Path to the index file
150     :return: A PackIndex loaded from the given path
151     """
152     f = GitFile(path, 'rb')
153     try:
154         return load_pack_index_file(path, f)
155     finally:
156         f.close()
157
158
159 def _load_file_contents(f, size=None):
160     fileno = getattr(f, 'fileno', None)
161     # Attempt to use mmap if possible
162     if fileno is not None:
163         fd = f.fileno()
164         if size is None:
165             size = os.fstat(fd).st_size
166         try:
167             contents = mmap.mmap(fd, size, access=mmap.ACCESS_READ)
168         except mmap.error:
169             # Perhaps a socket?
170             pass
171         else:
172             return contents, size
173     contents = f.read()
174     size = len(contents)
175     return contents, size
176
177
178 def load_pack_index_file(path, f):
179     """Load an index file from a file-like object.
180
181     :param path: Path for the index file
182     :param f: File-like object
183     :return: A PackIndex loaded from the given file
184     """
185     contents, size = _load_file_contents(f)
186     if contents[:4] == '\377tOc':
187         version = struct.unpack(">L", contents[4:8])[0]
188         if version == 2:
189             return PackIndex2(path, file=f, contents=contents,
190                 size=size)
191         else:
192             raise KeyError("Unknown pack index format %d" % version)
193     else:
194         return PackIndex1(path, file=f, contents=contents, size=size)
195
196
197 def bisect_find_sha(start, end, sha, unpack_name):
198     """Find a SHA in a data blob with sorted SHAs.
199
200     :param start: Start index of range to search
201     :param end: End index of range to search
202     :param sha: Sha to find
203     :param unpack_name: Callback to retrieve SHA by index
204     :return: Index of the SHA, or None if it wasn't found
205     """
206     assert start <= end
207     while start <= end:
208         i = (start + end)/2
209         file_sha = unpack_name(i)
210         x = cmp(file_sha, sha)
211         if x < 0:
212             start = i + 1
213         elif x > 0:
214             end = i - 1
215         else:
216             return i
217     return None
218
219
220 class PackIndex(object):
221     """An index in to a packfile.
222
223     Given a sha id of an object a pack index can tell you the location in the
224     packfile of that object if it has it.
225     """
226
227     def __eq__(self, other):
228         if not isinstance(other, PackIndex):
229             return False
230
231         for (name1, _, _), (name2, _, _) in izip(self.iterentries(),
232                                                  other.iterentries()):
233             if name1 != name2:
234                 return False
235         return True
236
237     def __ne__(self, other):
238         return not self.__eq__(other)
239
240     def __len__(self):
241         """Return the number of entries in this pack index."""
242         raise NotImplementedError(self.__len__)
243
244     def __iter__(self):
245         """Iterate over the SHAs in this pack."""
246         raise NotImplementedError(self.__iter__)
247
248     def iterentries(self):
249         """Iterate over the entries in this pack index.
250
251         :return: iterator over tuples with object name, offset in packfile and
252             crc32 checksum.
253         """
254         raise NotImplementedError(self.iterentries)
255
256     def get_pack_checksum(self):
257         """Return the SHA1 checksum stored for the corresponding packfile.
258
259         :return: 20-byte binary digest
260         """
261         raise NotImplementedError(self.get_pack_checksum)
262
263     def object_index(self, sha):
264         """Return the index in to the corresponding packfile for the object.
265
266         Given the name of an object it will return the offset that object
267         lives at within the corresponding pack file. If the pack file doesn't
268         have the object then None will be returned.
269         """
270         if len(sha) == 40:
271             sha = hex_to_sha(sha)
272         return self._object_index(sha)
273
274     def _object_index(self, sha):
275         """See object_index.
276
277         :param sha: A *binary* SHA string. (20 characters long)_
278         """
279         raise NotImplementedError(self._object_index)
280
281     def __iter__(self):
282         """Iterate over the SHAs in this pack."""
283         return imap(sha_to_hex, self._itersha())
284
285     def objects_sha1(self):
286         """Return the hex SHA1 over all the shas of all objects in this pack.
287
288         :note: This is used for the filename of the pack.
289         """
290         return iter_sha1(self._itersha())
291
292     def _itersha(self):
293         """Yield all the SHA1's of the objects in the index, sorted."""
294         raise NotImplementedError(self._itersha)
295
296
297 class MemoryPackIndex(PackIndex):
298     """Pack index that is stored entirely in memory."""
299
300     def __init__(self, entries, pack_checksum=None):
301         """Create a new MemoryPackIndex.
302
303         :param entries: Sequence of name, idx, crc32 (sorted)
304         :param pack_checksum: Optional pack checksum
305         """
306         self._by_sha = {}
307         for name, idx, crc32 in entries:
308             self._by_sha[name] = idx
309         self._entries = entries
310         self._pack_checksum = pack_checksum
311
312     def get_pack_checksum(self):
313         return self._pack_checksum
314
315     def __len__(self):
316         return len(self._entries)
317
318     def _object_index(self, sha):
319         return self._by_sha[sha][0]
320
321     def _itersha(self):
322         return iter(self._by_sha)
323
324     def iterentries(self):
325         return iter(self._entries)
326
327
328 class FilePackIndex(PackIndex):
329     """Pack index that is based on a file.
330
331     To do the loop it opens the file, and indexes first 256 4 byte groups
332     with the first byte of the sha id. The value in the four byte group indexed
333     is the end of the group that shares the same starting byte. Subtract one
334     from the starting byte and index again to find the start of the group.
335     The values are sorted by sha id within the group, so do the math to find
336     the start and end offset and then bisect in to find if the value is present.
337     """
338
339     def __init__(self, filename, file=None, contents=None, size=None):
340         """Create a pack index object.
341
342         Provide it with the name of the index file to consider, and it will map
343         it whenever required.
344         """
345         self._filename = filename
346         # Take the size now, so it can be checked each time we map the file to
347         # ensure that it hasn't changed.
348         if file is None:
349             self._file = GitFile(filename, 'rb')
350         else:
351             self._file = file
352         if contents is None:
353             self._contents, self._size = _load_file_contents(file, size)
354         else:
355             self._contents, self._size = (contents, size)
356
357     def __eq__(self, other):
358         # Quick optimization:
359         if (isinstance(other, FilePackIndex) and
360             self._fan_out_table != other._fan_out_table):
361             return False
362
363         return super(FilePackIndex, self).__eq__(other)
364
365     def close(self):
366         self._file.close()
367
368     def __len__(self):
369         """Return the number of entries in this pack index."""
370         return self._fan_out_table[-1]
371
372     def _unpack_entry(self, i):
373         """Unpack the i-th entry in the index file.
374
375         :return: Tuple with object name (SHA), offset in pack file and CRC32
376             checksum (if known).
377         """
378         raise NotImplementedError(self._unpack_entry)
379
380     def _unpack_name(self, i):
381         """Unpack the i-th name from the index file."""
382         raise NotImplementedError(self._unpack_name)
383
384     def _unpack_offset(self, i):
385         """Unpack the i-th object offset from the index file."""
386         raise NotImplementedError(self._unpack_offset)
387
388     def _unpack_crc32_checksum(self, i):
389         """Unpack the crc32 checksum for the i-th object from the index file."""
390         raise NotImplementedError(self._unpack_crc32_checksum)
391
392     def _itersha(self):
393         for i in range(len(self)):
394             yield self._unpack_name(i)
395
396     def iterentries(self):
397         """Iterate over the entries in this pack index.
398
399         :return: iterator over tuples with object name, offset in packfile and
400             crc32 checksum.
401         """
402         for i in range(len(self)):
403             yield self._unpack_entry(i)
404
405     def _read_fan_out_table(self, start_offset):
406         ret = []
407         for i in range(0x100):
408             fanout_entry = self._contents[start_offset+i*4:start_offset+(i+1)*4]
409             ret.append(struct.unpack(">L", fanout_entry)[0])
410         return ret
411
412     def check(self):
413         """Check that the stored checksum matches the actual checksum."""
414         actual = self.calculate_checksum()
415         stored = self.get_stored_checksum()
416         if actual != stored:
417             raise ChecksumMismatch(stored, actual)
418
419     def calculate_checksum(self):
420         """Calculate the SHA1 checksum over this pack index.
421
422         :return: This is a 20-byte binary digest
423         """
424         return make_sha(self._contents[:-20]).digest()
425
426     def get_pack_checksum(self):
427         """Return the SHA1 checksum stored for the corresponding packfile.
428
429         :return: 20-byte binary digest
430         """
431         return str(self._contents[-40:-20])
432
433     def get_stored_checksum(self):
434         """Return the SHA1 checksum stored for this index.
435
436         :return: 20-byte binary digest
437         """
438         return str(self._contents[-20:])
439
440     def _object_index(self, sha):
441         """See object_index.
442
443         :param sha: A *binary* SHA string. (20 characters long)_
444         """
445         assert len(sha) == 20
446         idx = ord(sha[0])
447         if idx == 0:
448             start = 0
449         else:
450             start = self._fan_out_table[idx-1]
451         end = self._fan_out_table[idx]
452         i = bisect_find_sha(start, end, sha, self._unpack_name)
453         if i is None:
454             raise KeyError(sha)
455         return self._unpack_offset(i)
456
457
458 class PackIndex1(FilePackIndex):
459     """Version 1 Pack Index file."""
460
461     def __init__(self, filename, file=None, contents=None, size=None):
462         super(PackIndex1, self).__init__(filename, file, contents, size)
463         self.version = 1
464         self._fan_out_table = self._read_fan_out_table(0)
465
466     def _unpack_entry(self, i):
467         (offset, name) = unpack_from(">L20s", self._contents,
468                                      (0x100 * 4) + (i * 24))
469         return (name, offset, None)
470
471     def _unpack_name(self, i):
472         offset = (0x100 * 4) + (i * 24) + 4
473         return self._contents[offset:offset+20]
474
475     def _unpack_offset(self, i):
476         offset = (0x100 * 4) + (i * 24)
477         return unpack_from(">L", self._contents, offset)[0]
478
479     def _unpack_crc32_checksum(self, i):
480         # Not stored in v1 index files
481         return None
482
483
484 class PackIndex2(FilePackIndex):
485     """Version 2 Pack Index file."""
486
487     def __init__(self, filename, file=None, contents=None, size=None):
488         super(PackIndex2, self).__init__(filename, file, contents, size)
489         assert self._contents[:4] == '\377tOc', "Not a v2 pack index file"
490         (self.version, ) = unpack_from(">L", self._contents, 4)
491         assert self.version == 2, "Version was %d" % self.version
492         self._fan_out_table = self._read_fan_out_table(8)
493         self._name_table_offset = 8 + 0x100 * 4
494         self._crc32_table_offset = self._name_table_offset + 20 * len(self)
495         self._pack_offset_table_offset = (self._crc32_table_offset +
496                                           4 * len(self))
497
498     def _unpack_entry(self, i):
499         return (self._unpack_name(i), self._unpack_offset(i),
500                 self._unpack_crc32_checksum(i))
501
502     def _unpack_name(self, i):
503         offset = self._name_table_offset + i * 20
504         return self._contents[offset:offset+20]
505
506     def _unpack_offset(self, i):
507         offset = self._pack_offset_table_offset + i * 4
508         return unpack_from(">L", self._contents, offset)[0]
509
510     def _unpack_crc32_checksum(self, i):
511         return unpack_from(">L", self._contents,
512                           self._crc32_table_offset + i * 4)[0]
513
514
515 def read_pack_header(read):
516     """Read the header of a pack file.
517
518     :param read: Read function
519     :return: Tuple with pack version and number of objects
520     """
521     header = read(12)
522     assert header[:4] == "PACK"
523     (version,) = unpack_from(">L", header, 4)
524     assert version in (2, 3), "Version was %d" % version
525     (num_objects,) = unpack_from(">L", header, 8)
526     return (version, num_objects)
527
528
529 def chunks_length(chunks):
530     return sum(imap(len, chunks))
531
532
533 def unpack_object(read_all, read_some=None):
534     """Unpack a Git object.
535
536     :param read_all: Read function that blocks until the number of requested
537         bytes are read.
538     :param read_some: Read function that returns at least one byte, but may not
539         return the number of bytes requested.
540     :return: tuple with type, uncompressed data, compressed size and tail data.
541     """
542     if read_some is None:
543         read_some = read_all
544     bytes = take_msb_bytes(read_all)
545     type = (bytes[0] >> 4) & 0x07
546     size = bytes[0] & 0x0f
547     for i, byte in enumerate(bytes[1:]):
548         size += (byte & 0x7f) << ((i * 7) + 4)
549     raw_base = len(bytes)
550     if type == OFS_DELTA:
551         bytes = take_msb_bytes(read_all)
552         raw_base += len(bytes)
553         assert not (bytes[-1] & 0x80)
554         delta_base_offset = bytes[0] & 0x7f
555         for byte in bytes[1:]:
556             delta_base_offset += 1
557             delta_base_offset <<= 7
558             delta_base_offset += (byte & 0x7f)
559         uncomp, comp_len, unused = read_zlib_chunks(read_some, size)
560         assert size == chunks_length(uncomp)
561         return type, (delta_base_offset, uncomp), comp_len+raw_base, unused
562     elif type == REF_DELTA:
563         basename = read_all(20)
564         raw_base += 20
565         uncomp, comp_len, unused = read_zlib_chunks(read_some, size)
566         assert size == chunks_length(uncomp)
567         return type, (basename, uncomp), comp_len+raw_base, unused
568     else:
569         uncomp, comp_len, unused = read_zlib_chunks(read_some, size)
570         assert chunks_length(uncomp) == size
571         return type, uncomp, comp_len+raw_base, unused
572
573
574 def _compute_object_size((num, obj)):
575     """Compute the size of a unresolved object for use with LRUSizeCache."""
576     if num in DELTA_TYPES:
577         return chunks_length(obj[1])
578     return chunks_length(obj)
579
580
581 class PackStreamReader(object):
582     """Class to read a pack stream.
583
584     The pack is read from a ReceivableProtocol using read() or recv() as
585     appropriate.
586     """
587
588     def __init__(self, read_all, read_some=None):
589         self.read_all = read_all
590         if read_some is None:
591             self.read_some = read_all
592         else:
593             self.read_some = read_some
594         self.sha = make_sha()
595         self._offset = 0
596         self._rbuf = StringIO()
597         # trailer is a deque to avoid memory allocation on small reads
598         self._trailer = deque()
599
600     def _read(self, read, size):
601         """Read up to size bytes using the given callback.
602
603         As a side effect, update the verifier's hash (excluding the last 20
604         bytes read) and write through to the output file.
605
606         :param read: The read callback to read from.
607         :param size: The maximum number of bytes to read; the particular
608             behavior is callback-specific.
609         """
610         data = read(size)
611
612         # maintain a trailer of the last 20 bytes we've read
613         n = len(data)
614         self._offset += n
615         tn = len(self._trailer)
616         if n >= 20:
617             to_pop = tn
618             to_add = 20
619         else:
620             to_pop = max(n + tn - 20, 0)
621             to_add = n
622         for _ in xrange(to_pop):
623             self.sha.update(self._trailer.popleft())
624         self._trailer.extend(data[-to_add:])
625
626         # hash everything but the trailer
627         self.sha.update(data[:-to_add])
628         return data
629
630     def _buf_len(self):
631         buf = self._rbuf
632         start = buf.tell()
633         buf.seek(0, SEEK_END)
634         end = buf.tell()
635         buf.seek(start)
636         return end - start
637
638     @property
639     def offset(self):
640         return self._offset - self._buf_len()
641
642     def read(self, size):
643         """Read, blocking until size bytes are read."""
644         buf_len = self._buf_len()
645         if buf_len >= size:
646             return self._rbuf.read(size)
647         buf_data = self._rbuf.read()
648         self._rbuf = StringIO()
649         return buf_data + self._read(self.read_all, size - buf_len)
650
651     def recv(self, size):
652         """Read up to size bytes, blocking until one byte is read."""
653         buf_len = self._buf_len()
654         if buf_len:
655             data = self._rbuf.read(size)
656             if size >= buf_len:
657                 self._rbuf = StringIO()
658             return data
659         return self._read(self.read_some, size)
660
661     def __len__(self):
662         return self._num_objects
663
664     def read_objects(self):
665         """Read the objects in this pack file.
666
667         :raise AssertionError: if there is an error in the pack format.
668         :raise ChecksumMismatch: if the checksum of the pack contents does not
669             match the checksum in the pack trailer.
670         :raise zlib.error: if an error occurred during zlib decompression.
671         :raise IOError: if an error occurred writing to the output file.
672         """
673         pack_version, self._num_objects = read_pack_header(self.read)
674         for i in xrange(self._num_objects):
675             type, uncomp, comp_len, unused = unpack_object(self.read, self.recv)
676             yield type, uncomp, comp_len
677
678             # prepend any unused data to current read buffer
679             buf = StringIO()
680             buf.write(unused)
681             buf.write(self._rbuf.read())
682             buf.seek(0)
683             self._rbuf = buf
684
685         pack_sha = sha_to_hex(''.join([c for c in self._trailer]))
686         calculated_sha = self.sha.hexdigest()
687         if pack_sha != calculated_sha:
688             raise ChecksumMismatch(pack_sha, calculated_sha)
689
690
691 class PackObjectIterator(object):
692
693     def __init__(self, pack, progress=None):
694         self.i = 0
695         self.offset = pack._header_size
696         self.num = len(pack)
697         self.map = pack._file
698         self._progress = progress
699
700     def __iter__(self):
701         return self
702
703     def __len__(self):
704         return self.num
705
706     def next(self):
707         if self.i == self.num:
708             raise StopIteration
709         self.map.seek(self.offset)
710         (type, obj, total_size, unused) = unpack_object(self.map.read)
711         self.map.seek(self.offset)
712         crc32 = zlib.crc32(self.map.read(total_size)) & 0xffffffff
713         ret = (self.offset, type, obj, crc32)
714         self.offset += total_size
715         if self._progress is not None:
716             self._progress(self.i, self.num)
717         self.i+=1
718         return ret
719
720 def obj_sha(type, chunks):
721     """Compute the SHA for a numeric type and object chunks."""
722     sha = make_sha()
723     sha.update(object_header(type, chunks_length(chunks)))
724     for chunk in chunks:
725         sha.update(chunk)
726     return sha.digest()
727
728
729 class PackData(object):
730     """The data contained in a packfile.
731
732     Pack files can be accessed both sequentially for exploding a pack, and
733     directly with the help of an index to retrieve a specific object.
734
735     The objects within are either complete or a delta aginst another.
736
737     The header is variable length. If the MSB of each byte is set then it
738     indicates that the subsequent byte is still part of the header.
739     For the first byte the next MS bits are the type, which tells you the type
740     of object, and whether it is a delta. The LS byte is the lowest bits of the
741     size. For each subsequent byte the LS 7 bits are the next MS bits of the
742     size, i.e. the last byte of the header contains the MS bits of the size.
743
744     For the complete objects the data is stored as zlib deflated data.
745     The size in the header is the uncompressed object size, so to uncompress
746     you need to just keep feeding data to zlib until you get an object back,
747     or it errors on bad data. This is done here by just giving the complete
748     buffer from the start of the deflated object on. This is bad, but until I
749     get mmap sorted out it will have to do.
750
751     Currently there are no integrity checks done. Also no attempt is made to
752     try and detect the delta case, or a request for an object at the wrong
753     position.  It will all just throw a zlib or KeyError.
754     """
755
756     def __init__(self, filename, file=None, size=None):
757         """Create a PackData object representing the pack in the given filename.
758
759         The file must exist and stay readable until the object is disposed of. It
760         must also stay the same size. It will be mapped whenever needed.
761
762         Currently there is a restriction on the size of the pack as the python
763         mmap implementation is flawed.
764         """
765         self._filename = filename
766         self._size = size
767         self._header_size = 12
768         if file is None:
769             self._file = GitFile(self._filename, 'rb')
770         else:
771             self._file = file
772         (version, self._num_objects) = read_pack_header(self._file.read)
773         self._offset_cache = LRUSizeCache(1024*1024*20,
774             compute_size=_compute_object_size)
775         self.pack = None
776
777     @classmethod
778     def from_file(cls, file, size):
779         return cls(str(file), file=file, size=size)
780
781     @classmethod
782     def from_path(cls, path):
783         return cls(filename=path)
784
785     def close(self):
786         self._file.close()
787
788     def _get_size(self):
789         if self._size is not None:
790             return self._size
791         self._size = os.path.getsize(self._filename)
792         if self._size < self._header_size:
793             errmsg = ("%s is too small for a packfile (%d < %d)" %
794                       (self._filename, self._size, self._header_size))
795             raise AssertionError(errmsg)
796         return self._size
797
798     def __len__(self):
799         """Returns the number of objects in this pack."""
800         return self._num_objects
801
802     def calculate_checksum(self):
803         """Calculate the checksum for this pack.
804
805         :return: 20-byte binary SHA1 digest
806         """
807         s = make_sha()
808         self._file.seek(0)
809         todo = self._get_size() - 20
810         while todo > 0:
811             x = self._file.read(min(todo, 1<<16))
812             s.update(x)
813             todo -= len(x)
814         return s.digest()
815
816     def get_ref(self, sha):
817         """Get the object for a ref SHA, only looking in this pack."""
818         # TODO: cache these results
819         if self.pack is None:
820             raise KeyError(sha)
821         offset = self.pack.index.object_index(sha)
822         if not offset:
823             raise KeyError(sha)
824         type, obj = self.get_object_at(offset)
825         return offset, type, obj
826
827     def resolve_object(self, offset, type, obj, get_ref=None):
828         """Resolve an object, possibly resolving deltas when necessary.
829
830         :return: Tuple with object type and contents.
831         """
832         if type not in DELTA_TYPES:
833             return type, obj
834
835         if get_ref is None:
836             get_ref = self.get_ref
837         if type == OFS_DELTA:
838             (delta_offset, delta) = obj
839             # TODO: clean up asserts and replace with nicer error messages
840             assert isinstance(offset, int)
841             assert isinstance(delta_offset, int)
842             base_offset = offset-delta_offset
843             type, base_obj = self.get_object_at(base_offset)
844             assert isinstance(type, int)
845         elif type == REF_DELTA:
846             (basename, delta) = obj
847             assert isinstance(basename, str) and len(basename) == 20
848             base_offset, type, base_obj = get_ref(basename)
849             assert isinstance(type, int)
850         type, base_chunks = self.resolve_object(base_offset, type, base_obj)
851         chunks = apply_delta(base_chunks, delta)
852         # TODO(dborowitz): This can result in poor performance if large base
853         # objects are separated from deltas in the pack. We should reorganize
854         # so that we apply deltas to all objects in a chain one after the other
855         # to optimize cache performance.
856         if offset is not None:
857             self._offset_cache[offset] = type, chunks
858         return type, chunks
859
860     def iterobjects(self, progress=None):
861         return PackObjectIterator(self, progress)
862
863     def iterentries(self, progress=None):
864         """Yield entries summarizing the contents of this pack.
865
866         :param progress: Progress function, called with current and total
867             object count.
868         :return: iterator of tuples with (sha, offset, crc32)
869         """
870         for offset, type, obj, crc32 in self.iterobjects(progress=progress):
871             assert isinstance(offset, int)
872             assert isinstance(type, int)
873             assert isinstance(obj, list) or isinstance(obj, tuple)
874             type, obj = self.resolve_object(offset, type, obj)
875             yield obj_sha(type, obj), offset, crc32
876
877     def sorted_entries(self, progress=None):
878         """Return entries in this pack, sorted by SHA.
879
880         :param progress: Progress function, called with current and total
881             object count
882         :return: List of tuples with (sha, offset, crc32)
883         """
884         ret = list(self.iterentries(progress=progress))
885         ret.sort()
886         return ret
887
888     def create_index_v1(self, filename, progress=None):
889         """Create a version 1 file for this data file.
890
891         :param filename: Index filename.
892         :param progress: Progress report function
893         :return: Checksum of index file
894         """
895         entries = self.sorted_entries(progress=progress)
896         f = GitFile(filename, 'wb')
897         try:
898             return write_pack_index_v1(f, entries, self.calculate_checksum())
899         finally:
900             f.close()
901
902     def create_index_v2(self, filename, progress=None):
903         """Create a version 2 index file for this data file.
904
905         :param filename: Index filename.
906         :param progress: Progress report function
907         :return: Checksum of index file
908         """
909         entries = self.sorted_entries(progress=progress)
910         f = GitFile(filename, 'wb')
911         try:
912             return write_pack_index_v2(f, entries, self.calculate_checksum())
913         finally:
914             f.close()
915
916     def create_index(self, filename, progress=None,
917                      version=2):
918         """Create an  index file for this data file.
919
920         :param filename: Index filename.
921         :param progress: Progress report function
922         :return: Checksum of index file
923         """
924         if version == 1:
925             return self.create_index_v1(filename, progress)
926         elif version == 2:
927             return self.create_index_v2(filename, progress)
928         else:
929             raise ValueError("unknown index format %d" % version)
930
931     def get_stored_checksum(self):
932         """Return the expected checksum stored in this pack."""
933         self._file.seek(self._get_size()-20)
934         return self._file.read(20)
935
936     def check(self):
937         """Check the consistency of this pack."""
938         actual = self.calculate_checksum()
939         stored = self.get_stored_checksum()
940         if actual != stored:
941             raise ChecksumMismatch(stored, actual)
942
943     def get_object_at(self, offset):
944         """Given an offset in to the packfile return the object that is there.
945
946         Using the associated index the location of an object can be looked up,
947         and then the packfile can be asked directly for that object using this
948         function.
949         """
950         if offset in self._offset_cache:
951             return self._offset_cache[offset]
952         assert isinstance(offset, long) or isinstance(offset, int),\
953                 "offset was %r" % offset
954         assert offset >= self._header_size
955         self._file.seek(offset)
956         return unpack_object(self._file.read)[:2]
957
958
959 class ThinPackData(PackData):
960     """PackData for thin packs, which require an ObjectStore for resolving."""
961
962     def __init__(self, resolve_ext_ref, *args, **kwargs):
963         super(ThinPackData, self).__init__(*args, **kwargs)
964         self.resolve_ext_ref = resolve_ext_ref
965
966     @classmethod
967     def from_file(cls, resolve_ext_ref, file, size):
968         return cls(resolve_ext_ref, str(file), file=file, size=size)
969
970     def get_ref(self, sha):
971         """Resolve a reference looking in both this pack and the store."""
972         try:
973             # As part of completing a pack we create a Pack object with a
974             # ThinPackData and a full PackIndex, so check in the index first if
975             # possible.
976             # TODO(dborowitz): reevaluate this when the pack completion code is
977             # rewritten.
978             return super(ThinPackData, self).get_ref(sha)
979         except KeyError:
980             type, obj = self.resolve_ext_ref(sha)
981             return None, type, obj
982
983     def iterentries(self, progress=None):
984         """Yield entries summarizing the contents of this pack.
985
986         :param progress: Progress function, called with current and
987             total object count.
988
989         This will yield tuples with (sha, offset, crc32)
990         """
991         found = {}
992         postponed = defaultdict(list)
993
994         class Postpone(Exception):
995             """Raised to postpone delta resolving."""
996
997             def __init__(self, sha):
998                 self.sha = sha
999
1000         def get_ref_text(sha):
1001             assert len(sha) == 20
1002             if sha in found:
1003                 offset = found[sha]
1004                 type, obj = self.get_object_at(offset)
1005                 return offset, type, obj
1006             try:
1007                 return self.get_ref(sha)
1008             except KeyError:
1009                 raise Postpone(sha)
1010
1011         extra = []
1012         todo = chain(self.iterobjects(progress=progress), extra)
1013         for (offset, type, obj, crc32) in todo:
1014             assert isinstance(offset, int)
1015             if obj is None:
1016                 # Inflate postponed delta
1017                 obj, type = self.get_object_at(offset)
1018             assert isinstance(type, int)
1019             assert isinstance(obj, list) or isinstance(obj, tuple)
1020             try:
1021                 type, obj = self.resolve_object(offset, type, obj, get_ref_text)
1022             except Postpone, e:
1023                 # Save memory by not storing the inflated obj in postponed
1024                 postponed[e.sha].append((offset, type, None, crc32))
1025             else:
1026                 sha = obj_sha(type, obj)
1027                 found[sha] = offset
1028                 yield sha, offset, crc32
1029                 extra.extend(postponed.pop(sha, []))
1030         if postponed:
1031             raise KeyError([sha_to_hex(h) for h in postponed.keys()])
1032
1033
1034 class SHA1Reader(object):
1035     """Wrapper around a file-like object that remembers the SHA1 of its data."""
1036
1037     def __init__(self, f):
1038         self.f = f
1039         self.sha1 = make_sha("")
1040
1041     def read(self, num=None):
1042         data = self.f.read(num)
1043         self.sha1.update(data)
1044         return data
1045
1046     def check_sha(self):
1047         stored = self.f.read(20)
1048         if stored != self.sha1.digest():
1049             raise ChecksumMismatch(self.sha1.hexdigest(), sha_to_hex(stored))
1050
1051     def close(self):
1052         return self.f.close()
1053
1054     def tell(self):
1055         return self.f.tell()
1056
1057
1058 class SHA1Writer(object):
1059     """Wrapper around a file-like object that remembers the SHA1 of its data."""
1060
1061     def __init__(self, f):
1062         self.f = f
1063         self.sha1 = make_sha("")
1064
1065     def write(self, data):
1066         self.sha1.update(data)
1067         self.f.write(data)
1068
1069     def write_sha(self):
1070         sha = self.sha1.digest()
1071         assert len(sha) == 20
1072         self.f.write(sha)
1073         return sha
1074
1075     def close(self):
1076         sha = self.write_sha()
1077         self.f.close()
1078         return sha
1079
1080     def tell(self):
1081         return self.f.tell()
1082
1083
1084 def write_pack_object(f, type, object):
1085     """Write pack object to a file.
1086
1087     :param f: File to write to
1088     :param type: Numeric type of the object
1089     :param object: Object to write
1090     :return: Tuple with offset at which the object was written, and crc32
1091     """
1092     offset = f.tell()
1093     packed_data_hdr = ""
1094     if type == OFS_DELTA:
1095         (delta_base_offset, object) = object
1096     elif type == REF_DELTA:
1097         (basename, object) = object
1098     size = len(object)
1099     c = (type << 4) | (size & 15)
1100     size >>= 4
1101     while size:
1102         packed_data_hdr += (chr(c | 0x80))
1103         c = size & 0x7f
1104         size >>= 7
1105     packed_data_hdr += chr(c)
1106     if type == OFS_DELTA:
1107         ret = [delta_base_offset & 0x7f]
1108         delta_base_offset >>= 7
1109         while delta_base_offset:
1110             delta_base_offset -= 1
1111             ret.insert(0, 0x80 | (delta_base_offset & 0x7f))
1112             delta_base_offset >>= 7
1113         packed_data_hdr += "".join([chr(x) for x in ret])
1114     elif type == REF_DELTA:
1115         assert len(basename) == 20
1116         packed_data_hdr += basename
1117     packed_data = packed_data_hdr + zlib.compress(object)
1118     f.write(packed_data)
1119     return (offset, (zlib.crc32(packed_data) & 0xffffffff))
1120
1121
1122 def write_pack(filename, objects, num_objects):
1123     """Write a new pack data file.
1124
1125     :param filename: Path to the new pack file (without .pack extension)
1126     :param objects: Iterable over (object, path) tuples to write
1127     :param num_objects: Number of objects to write
1128     :return: Tuple with checksum of pack file and index file
1129     """
1130     f = GitFile(filename + ".pack", 'wb')
1131     try:
1132         entries, data_sum = write_pack_data(f, objects, num_objects)
1133     finally:
1134         f.close()
1135     entries.sort()
1136     f = GitFile(filename + ".idx", 'wb')
1137     try:
1138         return data_sum, write_pack_index_v2(f, entries, data_sum)
1139     finally:
1140         f.close()
1141
1142
1143 def write_pack_header(f, num_objects):
1144     """Write a pack header for the given number of objects."""
1145     f.write('PACK')                          # Pack header
1146     f.write(struct.pack('>L', 2))            # Pack version
1147     f.write(struct.pack('>L', num_objects))  # Number of objects in pack
1148
1149
1150 def write_pack_data(f, objects, num_objects, window=10):
1151     """Write a new pack data file.
1152
1153     :param f: File to write to
1154     :param objects: Iterable over (object, path) tuples to write
1155     :param num_objects: Number of objects to write
1156     :param window: Sliding window size for searching for deltas; currently
1157                    unimplemented
1158     :return: List with (name, offset, crc32 checksum) entries, pack checksum
1159     """
1160     recency = list(objects)
1161     # FIXME: Somehow limit delta depth
1162     # FIXME: Make thin-pack optional (its not used when cloning a pack)
1163     # Build a list of objects ordered by the magic Linus heuristic
1164     # This helps us find good objects to diff against us
1165     magic = []
1166     for obj, path in recency:
1167         magic.append( (obj.type_num, path, 1, -obj.raw_length(), obj) )
1168     magic.sort()
1169     # Build a map of objects and their index in magic - so we can find
1170     # preceeding objects to diff against
1171     offs = {}
1172     for i in range(len(magic)):
1173         offs[magic[i][4]] = i
1174     # Write the pack
1175     entries = []
1176     f = SHA1Writer(f)
1177     write_pack_header(f, num_objects)
1178     for o, path in recency:
1179         sha1 = o.sha().digest()
1180         orig_t = o.type_num
1181         raw = o.as_raw_string()
1182         winner = raw
1183         t = orig_t
1184         #for i in range(offs[o]-window, window):
1185         #    if i < 0 or i >= len(offs): continue
1186         #    b = magic[i][4]
1187         #    if b.type_num != orig_t: continue
1188         #    base = b.as_raw_string()
1189         #    delta = create_delta(base, raw)
1190         #    if len(delta) < len(winner):
1191         #        winner = delta
1192         #        t = 6 if magic[i][2] == 1 else 7
1193         offset, crc32 = write_pack_object(f, t, winner)
1194         entries.append((sha1, offset, crc32))
1195     return entries, f.write_sha()
1196
1197
1198 def write_pack_index_v1(f, entries, pack_checksum):
1199     """Write a new pack index file.
1200
1201     :param f: A file-like object to write to
1202     :param entries: List of tuples with object name (sha), offset_in_pack,
1203         and crc32_checksum.
1204     :param pack_checksum: Checksum of the pack file.
1205     :return: The SHA of the written index file
1206     """
1207     f = SHA1Writer(f)
1208     fan_out_table = defaultdict(lambda: 0)
1209     for (name, offset, entry_checksum) in entries:
1210         fan_out_table[ord(name[0])] += 1
1211     # Fan-out table
1212     for i in range(0x100):
1213         f.write(struct.pack(">L", fan_out_table[i]))
1214         fan_out_table[i+1] += fan_out_table[i]
1215     for (name, offset, entry_checksum) in entries:
1216         f.write(struct.pack(">L20s", offset, name))
1217     assert len(pack_checksum) == 20
1218     f.write(pack_checksum)
1219     return f.write_sha()
1220
1221
1222 def create_delta(base_buf, target_buf):
1223     """Use python difflib to work out how to transform base_buf to target_buf.
1224
1225     :param base_buf: Base buffer
1226     :param target_buf: Target buffer
1227     """
1228     assert isinstance(base_buf, str)
1229     assert isinstance(target_buf, str)
1230     out_buf = ""
1231     # write delta header
1232     def encode_size(size):
1233         ret = ""
1234         c = size & 0x7f
1235         size >>= 7
1236         while size:
1237             ret += chr(c | 0x80)
1238             c = size & 0x7f
1239             size >>= 7
1240         ret += chr(c)
1241         return ret
1242     out_buf += encode_size(len(base_buf))
1243     out_buf += encode_size(len(target_buf))
1244     # write out delta opcodes
1245     seq = difflib.SequenceMatcher(a=base_buf, b=target_buf)
1246     for opcode, i1, i2, j1, j2 in seq.get_opcodes():
1247         # Git patch opcodes don't care about deletes!
1248         #if opcode == "replace" or opcode == "delete":
1249         #    pass
1250         if opcode == "equal":
1251             # If they are equal, unpacker will use data from base_buf
1252             # Write out an opcode that says what range to use
1253             scratch = ""
1254             op = 0x80
1255             o = i1
1256             for i in range(4):
1257                 if o & 0xff << i*8:
1258                     scratch += chr((o >> i*8) & 0xff)
1259                     op |= 1 << i
1260             s = i2 - i1
1261             for i in range(2):
1262                 if s & 0xff << i*8:
1263                     scratch += chr((s >> i*8) & 0xff)
1264                     op |= 1 << (4+i)
1265             out_buf += chr(op)
1266             out_buf += scratch
1267         if opcode == "replace" or opcode == "insert":
1268             # If we are replacing a range or adding one, then we just
1269             # output it to the stream (prefixed by its size)
1270             s = j2 - j1
1271             o = j1
1272             while s > 127:
1273                 out_buf += chr(127)
1274                 out_buf += target_buf[o:o+127]
1275                 s -= 127
1276                 o += 127
1277             out_buf += chr(s)
1278             out_buf += target_buf[o:o+s]
1279     return out_buf
1280
1281
1282 def apply_delta(src_buf, delta):
1283     """Based on the similar function in git's patch-delta.c.
1284
1285     :param src_buf: Source buffer
1286     :param delta: Delta instructions
1287     """
1288     if type(src_buf) != str:
1289         src_buf = "".join(src_buf)
1290     if type(delta) != str:
1291         delta = "".join(delta)
1292     out = []
1293     index = 0
1294     delta_length = len(delta)
1295     def get_delta_header_size(delta, index):
1296         size = 0
1297         i = 0
1298         while delta:
1299             cmd = ord(delta[index])
1300             index += 1
1301             size |= (cmd & ~0x80) << i
1302             i += 7
1303             if not cmd & 0x80:
1304                 break
1305         return size, index
1306     src_size, index = get_delta_header_size(delta, index)
1307     dest_size, index = get_delta_header_size(delta, index)
1308     assert src_size == len(src_buf), "%d vs %d" % (src_size, len(src_buf))
1309     while index < delta_length:
1310         cmd = ord(delta[index])
1311         index += 1
1312         if cmd & 0x80:
1313             cp_off = 0
1314             for i in range(4):
1315                 if cmd & (1 << i):
1316                     x = ord(delta[index])
1317                     index += 1
1318                     cp_off |= x << (i * 8)
1319             cp_size = 0
1320             for i in range(3):
1321                 if cmd & (1 << (4+i)):
1322                     x = ord(delta[index])
1323                     index += 1
1324                     cp_size |= x << (i * 8)
1325             if cp_size == 0:
1326                 cp_size = 0x10000
1327             if (cp_off + cp_size < cp_size or
1328                 cp_off + cp_size > src_size or
1329                 cp_size > dest_size):
1330                 break
1331             out.append(src_buf[cp_off:cp_off+cp_size])
1332         elif cmd != 0:
1333             out.append(delta[index:index+cmd])
1334             index += cmd
1335         else:
1336             raise ApplyDeltaError("Invalid opcode 0")
1337
1338     if index != delta_length:
1339         raise ApplyDeltaError("delta not empty: %r" % delta[index:])
1340
1341     if dest_size != chunks_length(out):
1342         raise ApplyDeltaError("dest size incorrect")
1343
1344     return out
1345
1346
1347 def write_pack_index_v2(f, entries, pack_checksum):
1348     """Write a new pack index file.
1349
1350     :param f: File-like object to write to
1351     :param entries: List of tuples with object name (sha), offset_in_pack, and
1352         crc32_checksum.
1353     :param pack_checksum: Checksum of the pack file.
1354     :return: The SHA of the index file written
1355     """
1356     f = SHA1Writer(f)
1357     f.write('\377tOc') # Magic!
1358     f.write(struct.pack(">L", 2))
1359     fan_out_table = defaultdict(lambda: 0)
1360     for (name, offset, entry_checksum) in entries:
1361         fan_out_table[ord(name[0])] += 1
1362     # Fan-out table
1363     for i in range(0x100):
1364         f.write(struct.pack(">L", fan_out_table[i]))
1365         fan_out_table[i+1] += fan_out_table[i]
1366     for (name, offset, entry_checksum) in entries:
1367         f.write(name)
1368     for (name, offset, entry_checksum) in entries:
1369         f.write(struct.pack(">L", entry_checksum))
1370     for (name, offset, entry_checksum) in entries:
1371         # FIXME: handle if MSBit is set in offset
1372         f.write(struct.pack(">L", offset))
1373     # FIXME: handle table for pack files > 8 Gb
1374     assert len(pack_checksum) == 20
1375     f.write(pack_checksum)
1376     return f.write_sha()
1377
1378
1379 class Pack(object):
1380     """A Git pack object."""
1381
1382     def __init__(self, basename):
1383         self._basename = basename
1384         self._data = None
1385         self._idx = None
1386         self._idx_path = self._basename + ".idx"
1387         self._data_path = self._basename + ".pack"
1388         self._data_load = lambda: PackData(self._data_path)
1389         self._idx_load = lambda: load_pack_index(self._idx_path)
1390
1391     @classmethod
1392     def from_lazy_objects(self, data_fn, idx_fn):
1393         """Create a new pack object from callables to load pack data and 
1394         index objects."""
1395         ret = Pack("")
1396         ret._data_load = data_fn
1397         ret._idx_load = idx_fn
1398         return ret
1399
1400     @classmethod
1401     def from_objects(self, data, idx):
1402         """Create a new pack object from pack data and index objects."""
1403         ret = Pack("")
1404         ret._data_load = lambda: data
1405         ret._idx_load = lambda: idx
1406         return ret
1407
1408     def name(self):
1409         """The SHA over the SHAs of the objects in this pack."""
1410         return self.index.objects_sha1()
1411
1412     @property
1413     def data(self):
1414         """The pack data object being used."""
1415         if self._data is None:
1416             self._data = self._data_load()
1417             self._data.pack = self
1418             assert len(self.index) == len(self._data)
1419             idx_stored_checksum = self.index.get_pack_checksum()
1420             data_stored_checksum = self._data.get_stored_checksum()
1421             if idx_stored_checksum != data_stored_checksum:
1422                 raise ChecksumMismatch(sha_to_hex(idx_stored_checksum),
1423                                        sha_to_hex(data_stored_checksum))
1424         return self._data
1425
1426     @property
1427     def index(self):
1428         """The index being used.
1429
1430         :note: This may be an in-memory index
1431         """
1432         if self._idx is None:
1433             self._idx = self._idx_load()
1434         return self._idx
1435
1436     def close(self):
1437         if self._data is not None:
1438             self._data.close()
1439         self.index.close()
1440
1441     def __eq__(self, other):
1442         return type(self) == type(other) and self.index == other.index
1443
1444     def __len__(self):
1445         """Number of entries in this pack."""
1446         return len(self.index)
1447
1448     def __repr__(self):
1449         return "%s(%r)" % (self.__class__.__name__, self._basename)
1450
1451     def __iter__(self):
1452         """Iterate over all the sha1s of the objects in this pack."""
1453         return iter(self.index)
1454
1455     def check(self):
1456         """Check the integrity of this pack.
1457
1458         :raise ChecksumMismatch: if a checksum for the index or data is wrong
1459         """
1460         self.index.check()
1461         self.data.check()
1462         for obj in self.iterobjects():
1463             obj.check()
1464         # TODO: object connectivity checks
1465
1466     def get_stored_checksum(self):
1467         return self.data.get_stored_checksum()
1468
1469     def __contains__(self, sha1):
1470         """Check whether this pack contains a particular SHA1."""
1471         try:
1472             self.index.object_index(sha1)
1473             return True
1474         except KeyError:
1475             return False
1476
1477     def get_raw(self, sha1):
1478         offset = self.index.object_index(sha1)
1479         obj_type, obj = self.data.get_object_at(offset)
1480         if type(offset) is long:
1481           offset = int(offset)
1482         type_num, chunks = self.data.resolve_object(offset, obj_type, obj)
1483         return type_num, "".join(chunks)
1484
1485     def __getitem__(self, sha1):
1486         """Retrieve the specified SHA1."""
1487         type, uncomp = self.get_raw(sha1)
1488         return ShaFile.from_raw_string(type, uncomp)
1489
1490     def iterobjects(self):
1491         """Iterate over the objects in this pack."""
1492         for offset, type, obj, crc32 in self.data.iterobjects():
1493             assert isinstance(offset, int)
1494             yield ShaFile.from_raw_chunks(
1495               *self.data.resolve_object(offset, type, obj))
1496
1497
1498 try:
1499     from dulwich._pack import apply_delta, bisect_find_sha
1500 except ImportError:
1501     pass