dulwich/object_store.py

   1 # object_store.py -- Object store for git objects
   2 # Copyright (C) 2008-2009 Jelmer Vernooij <jelmer@samba.org>
   3 #
   4 # This program is free software; you can redistribute it and/or
   5 # modify it under the terms of the GNU General Public License
   6 # as published by the Free Software Foundation; either version 2
   7 # or (at your option) a later version of the License.
   8 #
   9 # This program is distributed in the hope that it will be useful,
  10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 # GNU General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU General Public License
  15 # along with this program; if not, write to the Free Software
  16 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
  17 # MA  02110-1301, USA.
  18
  19
  20 """Git object store interfaces and implementation."""
  21
  22
  23 import errno
  24 import itertools
  25 import os
  26 import stat
  27 import tempfile
  28
  29 from dulwich.diff_tree import (
  30     tree_changes,
  31     walk_trees,
  32     )
  33 from dulwich.errors import (
  34     NotTreeError,
  35     )
  36 from dulwich.file import GitFile
  37 from dulwich.objects import (
  38     Commit,
  39     ShaFile,
  40     Tag,
  41     Tree,
  42     ZERO_SHA,
  43     hex_to_sha,
  44     sha_to_hex,
  45     hex_to_filename,
  46     S_ISGITLINK,
  47     object_class,
  48     )
  49 from dulwich.pack import (
  50     Pack,
  51     PackData,
  52     iter_sha1,
  53     write_pack_header,
  54     write_pack_index_v2,
  55     write_pack_object,
  56     write_pack_objects,
  57     compute_file_sha,
  58     PackIndexer,
  59     PackStreamCopier,
  60     )
  61
  62 INFODIR = 'info'
  63 PACKDIR = 'pack'
  64
  65
  66 class BaseObjectStore(object):
  67     """Object store interface."""
  68
  69     def determine_wants_all(self, refs):
  70         return [sha for (ref, sha) in refs.iteritems()
  71                 if not sha in self and not ref.endswith("^{}") and
  72                    not sha == ZERO_SHA]
  73
  74     def iter_shas(self, shas):
  75         """Iterate over the objects for the specified shas.
  76
  77         :param shas: Iterable object with SHAs
  78         :return: Object iterator
  79         """
  80         return ObjectStoreIterator(self, shas)
  81
  82     def contains_loose(self, sha):
  83         """Check if a particular object is present by SHA1 and is loose."""
  84         raise NotImplementedError(self.contains_loose)
  85
  86     def contains_packed(self, sha):
  87         """Check if a particular object is present by SHA1 and is packed."""
  88         raise NotImplementedError(self.contains_packed)
  89
  90     def __contains__(self, sha):
  91         """Check if a particular object is present by SHA1.
  92
  93         This method makes no distinction between loose and packed objects.
  94         """
  95         return self.contains_packed(sha) or self.contains_loose(sha)
  96
  97     @property
  98     def packs(self):
  99         """Iterable of pack objects."""
 100         raise NotImplementedError
 101
 102     def get_raw(self, name):
 103         """Obtain the raw text for an object.
 104
 105         :param name: sha for the object.
 106         :return: tuple with numeric type and object contents.
 107         """
 108         raise NotImplementedError(self.get_raw)
 109
 110     def __getitem__(self, sha):
 111         """Obtain an object by SHA1."""
 112         type_num, uncomp = self.get_raw(sha)
 113         return ShaFile.from_raw_string(type_num, uncomp)
 114
 115     def __iter__(self):
 116         """Iterate over the SHAs that are present in this store."""
 117         raise NotImplementedError(self.__iter__)
 118
 119     def add_object(self, obj):
 120         """Add a single object to this object store.
 121
 122         """
 123         raise NotImplementedError(self.add_object)
 124
 125     def add_objects(self, objects):
 126         """Add a set of objects to this object store.
 127
 128         :param objects: Iterable over a list of objects.
 129         """
 130         raise NotImplementedError(self.add_objects)
 131
 132     def tree_changes(self, source, target, want_unchanged=False):
 133         """Find the differences between the contents of two trees
 134
 135         :param object_store: Object store to use for retrieving tree contents
 136         :param tree: SHA1 of the root tree
 137         :param want_unchanged: Whether unchanged files should be reported
 138         :return: Iterator over tuples with
 139             (oldpath, newpath), (oldmode, newmode), (oldsha, newsha)
 140         """
 141         for change in tree_changes(self, source, target,
 142                                    want_unchanged=want_unchanged):
 143             yield ((change.old.path, change.new.path),
 144                    (change.old.mode, change.new.mode),
 145                    (change.old.sha, change.new.sha))
 146
 147     def iter_tree_contents(self, tree_id, include_trees=False):
 148         """Iterate the contents of a tree and all subtrees.
 149
 150         Iteration is depth-first pre-order, as in e.g. os.walk.
 151
 152         :param tree_id: SHA1 of the tree.
 153         :param include_trees: If True, include tree objects in the iteration.
 154         :return: Iterator over TreeEntry namedtuples for all the objects in a
 155             tree.
 156         """
 157         for entry, _ in walk_trees(self, tree_id, None):
 158             if not stat.S_ISDIR(entry.mode) or include_trees:
 159                 yield entry
 160
 161     def find_missing_objects(self, haves, wants, progress=None,
 162                              get_tagged=None):
 163         """Find the missing objects required for a set of revisions.
 164
 165         :param haves: Iterable over SHAs already in common.
 166         :param wants: Iterable over SHAs of objects to fetch.
 167         :param progress: Simple progress function that will be called with
 168             updated progress strings.
 169         :param get_tagged: Function that returns a dict of pointed-to sha -> tag
 170             sha for including tags.
 171         :return: Iterator over (sha, path) pairs.
 172         """
 173         finder = MissingObjectFinder(self, haves, wants, progress, get_tagged)
 174         return iter(finder.next, None)
 175
 176     def find_common_revisions(self, graphwalker):
 177         """Find which revisions this store has in common using graphwalker.
 178
 179         :param graphwalker: A graphwalker object.
 180         :return: List of SHAs that are in common
 181         """
 182         haves = []
 183         sha = graphwalker.next()
 184         while sha:
 185             if sha in self:
 186                 haves.append(sha)
 187                 graphwalker.ack(sha)
 188             sha = graphwalker.next()
 189         return haves
 190
 191     def get_graph_walker(self, heads):
 192         """Obtain a graph walker for this object store.
 193
 194         :param heads: Local heads to start search with
 195         :return: GraphWalker object
 196         """
 197         return ObjectStoreGraphWalker(heads, lambda sha: self[sha].parents)
 198
 199     def generate_pack_contents(self, have, want, progress=None):
 200         """Iterate over the contents of a pack file.
 201
 202         :param have: List of SHA1s of objects that should not be sent
 203         :param want: List of SHA1s of objects that should be sent
 204         :param progress: Optional progress reporting method
 205         """
 206         return self.iter_shas(self.find_missing_objects(have, want, progress))
 207
 208     def peel_sha(self, sha):
 209         """Peel all tags from a SHA.
 210
 211         :param sha: The object SHA to peel.
 212         :return: The fully-peeled SHA1 of a tag object, after peeling all
 213             intermediate tags; if the original ref does not point to a tag, this
 214             will equal the original SHA1.
 215         """
 216         obj = self[sha]
 217         obj_class = object_class(obj.type_name)
 218         while obj_class is Tag:
 219             obj_class, sha = obj.object
 220             obj = self[sha]
 221         return obj
 222
 223
 224 class PackBasedObjectStore(BaseObjectStore):
 225
 226     def __init__(self):
 227         self._pack_cache = None
 228
 229     @property
 230     def alternates(self):
 231         return []
 232
 233     def contains_packed(self, sha):
 234         """Check if a particular object is present by SHA1 and is packed."""
 235         for pack in self.packs:
 236             if sha in pack:
 237                 return True
 238         return False
 239
 240     def _load_packs(self):
 241         raise NotImplementedError(self._load_packs)
 242
 243     def _pack_cache_stale(self):
 244         """Check whether the pack cache is stale."""
 245         raise NotImplementedError(self._pack_cache_stale)
 246
 247     def _add_known_pack(self, pack):
 248         """Add a newly appeared pack to the cache by path.
 249
 250         """
 251         if self._pack_cache is not None:
 252             self._pack_cache.append(pack)
 253
 254     @property
 255     def packs(self):
 256         """List with pack objects."""
 257         if self._pack_cache is None or self._pack_cache_stale():
 258             self._pack_cache = self._load_packs()
 259         return self._pack_cache
 260
 261     def _iter_loose_objects(self):
 262         """Iterate over the SHAs of all loose objects."""
 263         raise NotImplementedError(self._iter_loose_objects)
 264
 265     def _get_loose_object(self, sha):
 266         raise NotImplementedError(self._get_loose_object)
 267
 268     def _remove_loose_object(self, sha):
 269         raise NotImplementedError(self._remove_loose_object)
 270
 271     def pack_loose_objects(self):
 272         """Pack loose objects.
 273
 274         :return: Number of objects packed
 275         """
 276         objects = set()
 277         for sha in self._iter_loose_objects():
 278             objects.add((self._get_loose_object(sha), None))
 279         self.add_objects(list(objects))
 280         for obj, path in objects:
 281             self._remove_loose_object(obj.id)
 282         return len(objects)
 283
 284     def __iter__(self):
 285         """Iterate over the SHAs that are present in this store."""
 286         iterables = self.packs + [self._iter_loose_objects()]
 287         return itertools.chain(*iterables)
 288
 289     def contains_loose(self, sha):
 290         """Check if a particular object is present by SHA1 and is loose."""
 291         return self._get_loose_object(sha) is not None
 292
 293     def get_raw(self, name):
 294         """Obtain the raw text for an object.
 295
 296         :param name: sha for the object.
 297         :return: tuple with numeric type and object contents.
 298         """
 299         if len(name) == 40:
 300             sha = hex_to_sha(name)
 301             hexsha = name
 302         elif len(name) == 20:
 303             sha = name
 304             hexsha = None
 305         else:
 306             raise AssertionError("Invalid object name %r" % name)
 307         for pack in self.packs:
 308             try:
 309                 return pack.get_raw(sha)
 310             except KeyError:
 311                 pass
 312         if hexsha is None:
 313             hexsha = sha_to_hex(name)
 314         ret = self._get_loose_object(hexsha)
 315         if ret is not None:
 316             return ret.type_num, ret.as_raw_string()
 317         for alternate in self.alternates:
 318             try:
 319                 return alternate.get_raw(hexsha)
 320             except KeyError:
 321                 pass
 322         raise KeyError(hexsha)
 323
 324     def add_objects(self, objects):
 325         """Add a set of objects to this object store.
 326
 327         :param objects: Iterable over objects, should support __len__.
 328         :return: Pack object of the objects written.
 329         """
 330         if len(objects) == 0:
 331             # Don't bother writing an empty pack file
 332             return
 333         f, commit = self.add_pack()
 334         write_pack_objects(f, objects)
 335         return commit()
 336
 337
 338 class DiskObjectStore(PackBasedObjectStore):
 339     """Git-style object store that exists on disk."""
 340
 341     def __init__(self, path):
 342         """Open an object store.
 343
 344         :param path: Path of the object store.
 345         """
 346         super(DiskObjectStore, self).__init__()
 347         self.path = path
 348         self.pack_dir = os.path.join(self.path, PACKDIR)
 349         self._pack_cache_time = 0
 350         self._alternates = None
 351
 352     @property
 353     def alternates(self):
 354         if self._alternates is not None:
 355             return self._alternates
 356         self._alternates = []
 357         for path in self._read_alternate_paths():
 358             self._alternates.append(DiskObjectStore(path))
 359         return self._alternates
 360
 361     def _read_alternate_paths(self):
 362         try:
 363             f = GitFile(os.path.join(self.path, "info", "alternates"),
 364                     'rb')
 365         except (OSError, IOError), e:
 366             if e.errno == errno.ENOENT:
 367                 return []
 368             raise
 369         ret = []
 370         try:
 371             for l in f.readlines():
 372                 l = l.rstrip("\n")
 373                 if l[0] == "#":
 374                     continue
 375                 if not os.path.isabs(l):
 376                     continue
 377                 ret.append(l)
 378             return ret
 379         finally:
 380             f.close()
 381
 382     def add_alternate_path(self, path):
 383         """Add an alternate path to this object store.
 384         """
 385         try:
 386             os.mkdir(os.path.join(self.path, "info"))
 387         except OSError, e:
 388             if e.errno != errno.EEXIST:
 389                 raise
 390         alternates_path = os.path.join(self.path, "info/alternates")
 391         f = GitFile(alternates_path, 'wb')
 392         try:
 393             try:
 394                 orig_f = open(alternates_path, 'rb')
 395             except (OSError, IOError), e:
 396                 if e.errno != errno.ENOENT:
 397                     raise
 398             else:
 399                 try:
 400                     f.write(orig_f.read())
 401                 finally:
 402                     orig_f.close()
 403             f.write("%s\n" % path)
 404         finally:
 405             f.close()
 406         self.alternates.append(DiskObjectStore(path))
 407
 408     def _load_packs(self):
 409         pack_files = []
 410         try:
 411             self._pack_cache_time = os.stat(self.pack_dir).st_mtime
 412             pack_dir_contents = os.listdir(self.pack_dir)
 413             for name in pack_dir_contents:
 414                 # TODO: verify that idx exists first
 415                 if name.startswith("pack-") and name.endswith(".pack"):
 416                     filename = os.path.join(self.pack_dir, name)
 417                     pack_files.append((os.stat(filename).st_mtime, filename))
 418         except OSError, e:
 419             if e.errno == errno.ENOENT:
 420                 return []
 421             raise
 422         pack_files.sort(reverse=True)
 423         suffix_len = len(".pack")
 424         return [Pack(f[:-suffix_len]) for _, f in pack_files]
 425
 426     def _pack_cache_stale(self):
 427         try:
 428             return os.stat(self.pack_dir).st_mtime > self._pack_cache_time
 429         except OSError, e:
 430             if e.errno == errno.ENOENT:
 431                 return True
 432             raise
 433
 434     def _get_shafile_path(self, sha):
 435         # Check from object dir
 436         return hex_to_filename(self.path, sha)
 437
 438     def _iter_loose_objects(self):
 439         for base in os.listdir(self.path):
 440             if len(base) != 2:
 441                 continue
 442             for rest in os.listdir(os.path.join(self.path, base)):
 443                 yield base+rest
 444
 445     def _get_loose_object(self, sha):
 446         path = self._get_shafile_path(sha)
 447         try:
 448             return ShaFile.from_path(path)
 449         except (OSError, IOError), e:
 450             if e.errno == errno.ENOENT:
 451                 return None
 452             raise
 453
 454     def _remove_loose_object(self, sha):
 455         os.remove(self._get_shafile_path(sha))
 456
 457     def _complete_thin_pack(self, f, path, copier, indexer):
 458         """Move a specific file containing a pack into the pack directory.
 459
 460         :note: The file should be on the same file system as the
 461             packs directory.
 462
 463         :param f: Open file object for the pack.
 464         :param path: Path to the pack file.
 465         :param copier: A PackStreamCopier to use for writing pack data.
 466         :param indexer: A PackIndexer for indexing the pack.
 467         """
 468         entries = list(indexer)
 469
 470         # Update the header with the new number of objects.
 471         f.seek(0)
 472         write_pack_header(f, len(entries) + len(indexer.ext_refs()))
 473
 474         # Rescan the rest of the pack, computing the SHA with the new header.
 475         new_sha = compute_file_sha(f, end_ofs=-20)
 476
 477         # Complete the pack.
 478         for ext_sha in indexer.ext_refs():
 479             type_num, data = self.get_raw(ext_sha)
 480             offset = f.tell()
 481             crc32 = write_pack_object(f, type_num, data, sha=new_sha)
 482             entries.append((ext_sha, offset, crc32))
 483         pack_sha = new_sha.digest()
 484         f.write(pack_sha)
 485         f.close()
 486
 487         # Move the pack in.
 488         entries.sort()
 489         pack_base_name = os.path.join(
 490           self.pack_dir, 'pack-' + iter_sha1(e[0] for e in entries))
 491         os.rename(path, pack_base_name + '.pack')
 492
 493         # Write the index.
 494         index_file = GitFile(pack_base_name + '.idx', 'wb')
 495         try:
 496             write_pack_index_v2(index_file, entries, pack_sha)
 497             index_file.close()
 498         finally:
 499             index_file.abort()
 500
 501         # Add the pack to the store and return it.
 502         final_pack = Pack(pack_base_name)
 503         final_pack.check_length_and_checksum()
 504         self._add_known_pack(final_pack)
 505         return final_pack
 506
 507     def add_thin_pack(self, read_all, read_some):
 508         """Add a new thin pack to this object store.
 509
 510         Thin packs are packs that contain deltas with parents that exist outside
 511         the pack. They should never be placed in the object store directly, and
 512         always indexed and completed as they are copied.
 513
 514         :param read_all: Read function that blocks until the number of requested
 515             bytes are read.
 516         :param read_some: Read function that returns at least one byte, but may
 517             not return the number of bytes requested.
 518         :return: A Pack object pointing at the now-completed thin pack in the
 519             objects/pack directory.
 520         """
 521         fd, path = tempfile.mkstemp(dir=self.path, prefix='tmp_pack_')
 522         f = os.fdopen(fd, 'w+b')
 523
 524         try:
 525             indexer = PackIndexer(f, resolve_ext_ref=self.get_raw)
 526             copier = PackStreamCopier(read_all, read_some, f,
 527                                       delta_iter=indexer)
 528             copier.verify()
 529             return self._complete_thin_pack(f, path, copier, indexer)
 530         finally:
 531             f.close()
 532
 533     def move_in_pack(self, path):
 534         """Move a specific file containing a pack into the pack directory.
 535
 536         :note: The file should be on the same file system as the
 537             packs directory.
 538
 539         :param path: Path to the pack file.
 540         """
 541         p = PackData(path)
 542         entries = p.sorted_entries()
 543         basename = os.path.join(self.pack_dir,
 544             "pack-%s" % iter_sha1(entry[0] for entry in entries))
 545         f = GitFile(basename+".idx", "wb")
 546         try:
 547             write_pack_index_v2(f, entries, p.get_stored_checksum())
 548         finally:
 549             f.close()
 550         p.close()
 551         os.rename(path, basename + ".pack")
 552         final_pack = Pack(basename)
 553         self._add_known_pack(final_pack)
 554         return final_pack
 555
 556     def add_pack(self):
 557         """Add a new pack to this object store.
 558
 559         :return: Fileobject to write to and a commit function to
 560             call when the pack is finished.
 561         """
 562         fd, path = tempfile.mkstemp(dir=self.pack_dir, suffix=".pack")
 563         f = os.fdopen(fd, 'wb')
 564         def commit():
 565             os.fsync(fd)
 566             f.close()
 567             if os.path.getsize(path) > 0:
 568                 return self.move_in_pack(path)
 569             else:
 570                 os.remove(path)
 571                 return None
 572         return f, commit
 573
 574     def add_object(self, obj):
 575         """Add a single object to this object store.
 576
 577         :param obj: Object to add
 578         """
 579         dir = os.path.join(self.path, obj.id[:2])
 580         try:
 581             os.mkdir(dir)
 582         except OSError, e:
 583             if e.errno != errno.EEXIST:
 584                 raise
 585         path = os.path.join(dir, obj.id[2:])
 586         if os.path.exists(path):
 587             return # Already there, no need to write again
 588         f = GitFile(path, 'wb')
 589         try:
 590             f.write(obj.as_legacy_object())
 591         finally:
 592             f.close()
 593
 594     @classmethod
 595     def init(cls, path):
 596         try:
 597             os.mkdir(path)
 598         except OSError, e:
 599             if e.errno != errno.EEXIST:
 600                 raise
 601         os.mkdir(os.path.join(path, "info"))
 602         os.mkdir(os.path.join(path, PACKDIR))
 603         return cls(path)
 604
 605
 606 class MemoryObjectStore(BaseObjectStore):
 607     """Object store that keeps all objects in memory."""
 608
 609     def __init__(self):
 610         super(MemoryObjectStore, self).__init__()
 611         self._data = {}
 612
 613     def contains_loose(self, sha):
 614         """Check if a particular object is present by SHA1 and is loose."""
 615         return sha in self._data
 616
 617     def contains_packed(self, sha):
 618         """Check if a particular object is present by SHA1 and is packed."""
 619         return False
 620
 621     def __iter__(self):
 622         """Iterate over the SHAs that are present in this store."""
 623         return self._data.iterkeys()
 624
 625     @property
 626     def packs(self):
 627         """List with pack objects."""
 628         return []
 629
 630     def get_raw(self, name):
 631         """Obtain the raw text for an object.
 632
 633         :param name: sha for the object.
 634         :return: tuple with numeric type and object contents.
 635         """
 636         obj = self[name]
 637         return obj.type_num, obj.as_raw_string()
 638
 639     def __getitem__(self, name):
 640         return self._data[name]
 641
 642     def __delitem__(self, name):
 643         """Delete an object from this store, for testing only."""
 644         del self._data[name]
 645
 646     def add_object(self, obj):
 647         """Add a single object to this object store.
 648
 649         """
 650         self._data[obj.id] = obj
 651
 652     def add_objects(self, objects):
 653         """Add a set of objects to this object store.
 654
 655         :param objects: Iterable over a list of objects.
 656         """
 657         for obj, path in objects:
 658             self._data[obj.id] = obj
 659
 660
 661 class ObjectImporter(object):
 662     """Interface for importing objects."""
 663
 664     def __init__(self, count):
 665         """Create a new ObjectImporter.
 666
 667         :param count: Number of objects that's going to be imported.
 668         """
 669         self.count = count
 670
 671     def add_object(self, object):
 672         """Add an object."""
 673         raise NotImplementedError(self.add_object)
 674
 675     def finish(self, object):
 676         """Finish the import and write objects to disk."""
 677         raise NotImplementedError(self.finish)
 678
 679
 680 class ObjectIterator(object):
 681     """Interface for iterating over objects."""
 682
 683     def iterobjects(self):
 684         raise NotImplementedError(self.iterobjects)
 685
 686
 687 class ObjectStoreIterator(ObjectIterator):
 688     """ObjectIterator that works on top of an ObjectStore."""
 689
 690     def __init__(self, store, sha_iter):
 691         """Create a new ObjectIterator.
 692
 693         :param store: Object store to retrieve from
 694         :param sha_iter: Iterator over (sha, path) tuples
 695         """
 696         self.store = store
 697         self.sha_iter = sha_iter
 698         self._shas = []
 699
 700     def __iter__(self):
 701         """Yield tuple with next object and path."""
 702         for sha, path in self.itershas():
 703             yield self.store[sha], path
 704
 705     def iterobjects(self):
 706         """Iterate over just the objects."""
 707         for o, path in self:
 708             yield o
 709
 710     def itershas(self):
 711         """Iterate over the SHAs."""
 712         for sha in self._shas:
 713             yield sha
 714         for sha in self.sha_iter:
 715             self._shas.append(sha)
 716             yield sha
 717
 718     def __contains__(self, needle):
 719         """Check if an object is present.
 720
 721         :note: This checks if the object is present in
 722             the underlying object store, not if it would
 723             be yielded by the iterator.
 724
 725         :param needle: SHA1 of the object to check for
 726         """
 727         return needle in self.store
 728
 729     def __getitem__(self, key):
 730         """Find an object by SHA1.
 731
 732         :note: This retrieves the object from the underlying
 733             object store. It will also succeed if the object would
 734             not be returned by the iterator.
 735         """
 736         return self.store[key]
 737
 738     def __len__(self):
 739         """Return the number of objects."""
 740         return len(list(self.itershas()))
 741
 742
 743 def tree_lookup_path(lookup_obj, root_sha, path):
 744     """Look up an object in a Git tree.
 745
 746     :param lookup_obj: Callback for retrieving object by SHA1
 747     :param root_sha: SHA1 of the root tree
 748     :param path: Path to lookup
 749     :return: A tuple of (mode, SHA) of the resulting path.
 750     """
 751     tree = lookup_obj(root_sha)
 752     if not isinstance(tree, Tree):
 753         raise NotTreeError(root_sha)
 754     return tree.lookup_path(lookup_obj, path)
 755
 756
 757 class MissingObjectFinder(object):
 758     """Find the objects missing from another object store.
 759
 760     :param object_store: Object store containing at least all objects to be
 761         sent
 762     :param haves: SHA1s of commits not to send (already present in target)
 763     :param wants: SHA1s of commits to send
 764     :param progress: Optional function to report progress to.
 765     :param get_tagged: Function that returns a dict of pointed-to sha -> tag
 766         sha for including tags.
 767     :param tagged: dict of pointed-to sha -> tag sha for including tags
 768     """
 769
 770     def __init__(self, object_store, haves, wants, progress=None,
 771                  get_tagged=None):
 772         haves = set(haves)
 773         self.sha_done = haves
 774         self.objects_to_send = set([(w, None, False) for w in wants
 775                                     if w not in haves])
 776         self.object_store = object_store
 777         if progress is None:
 778             self.progress = lambda x: None
 779         else:
 780             self.progress = progress
 781         self._tagged = get_tagged and get_tagged() or {}
 782
 783     def add_todo(self, entries):
 784         self.objects_to_send.update([e for e in entries
 785                                      if not e[0] in self.sha_done])
 786
 787     def parse_tree(self, tree):
 788         self.add_todo([(sha, name, not stat.S_ISDIR(mode))
 789                        for name, mode, sha in tree.iteritems()
 790                        if not S_ISGITLINK(mode)])
 791
 792     def parse_commit(self, commit):
 793         self.add_todo([(commit.tree, "", False)])
 794         self.add_todo([(p, None, False) for p in commit.parents])
 795
 796     def parse_tag(self, tag):
 797         self.add_todo([(tag.object[1], None, False)])
 798
 799     def next(self):
 800         while True:
 801             if not self.objects_to_send:
 802                 return None
 803             (sha, name, leaf) = self.objects_to_send.pop()
 804             if sha not in self.sha_done:
 805                 break
 806         if not leaf:
 807             o = self.object_store[sha]
 808             if isinstance(o, Commit):
 809                 self.parse_commit(o)
 810             elif isinstance(o, Tree):
 811                 self.parse_tree(o)
 812             elif isinstance(o, Tag):
 813                 self.parse_tag(o)
 814         if sha in self._tagged:
 815             self.add_todo([(self._tagged[sha], None, True)])
 816         self.sha_done.add(sha)
 817         self.progress("counting objects: %d\r" % len(self.sha_done))
 818         return (sha, name)
 819
 820
 821 class ObjectStoreGraphWalker(object):
 822     """Graph walker that finds what commits are missing from an object store.
 823
 824     :ivar heads: Revisions without descendants in the local repo
 825     :ivar get_parents: Function to retrieve parents in the local repo
 826     """
 827
 828     def __init__(self, local_heads, get_parents):
 829         """Create a new instance.
 830
 831         :param local_heads: Heads to start search with
 832         :param get_parents: Function for finding the parents of a SHA1.
 833         """
 834         self.heads = set(local_heads)
 835         self.get_parents = get_parents
 836         self.parents = {}
 837
 838     def ack(self, sha):
 839         """Ack that a revision and its ancestors are present in the source."""
 840         ancestors = set([sha])
 841
 842         # stop if we run out of heads to remove
 843         while self.heads:
 844             for a in ancestors:
 845                 if a in self.heads:
 846                     self.heads.remove(a)
 847
 848             # collect all ancestors
 849             new_ancestors = set()
 850             for a in ancestors:
 851                 ps = self.parents.get(a)
 852                 if ps is not None:
 853                     new_ancestors.update(ps)
 854                 self.parents[a] = None
 855
 856             # no more ancestors; stop
 857             if not new_ancestors:
 858                 break
 859
 860             ancestors = new_ancestors
 861
 862     def next(self):
 863         """Iterate over ancestors of heads in the target."""
 864         if self.heads:
 865             ret = self.heads.pop()
 866             ps = self.get_parents(ret)
 867             self.parents[ret] = ps
 868             self.heads.update([p for p in ps if not p in self.parents])
 869             return ret
 870         return None