Refactor GraphWalker.

bin/dulwich

@@ -31,7 +31,6 @@ def get_transport_and_path(uri):
 
 
 def cmd_fetch_pack(args):
-	from dulwich.client import SimpleFetchGraphWalker
 	from dulwich.repo import Repo
 	opts, args = getopt(args, "", ["all"])
 	opts = dict(opts)
@@ -41,7 +40,7 @@ def cmd_fetch_pack(args):
 	else:
 		determine_wants = lambda x: [y for y in args if not y in r.object_store]
 	r = Repo(".")
-	graphwalker = SimpleFetchGraphWalker(r.heads().values(), r.get_parents)
+	graphwalker = r.get_graph_walker()
 	f, commit = r.object_store.add_pack()
 	try:
 		client.fetch_pack(path, determine_wants, graphwalker, f.write, sys.stdout.write)
@@ -143,7 +142,6 @@ def cmd_init(args):
 
 
 def cmd_clone(args):
-	from dulwich.client import SimpleFetchGraphWalker
 	from dulwich.repo import Repo
 	import os
 	import sys
@@ -164,7 +162,7 @@ def cmd_clone(args):
 		os.mkdir(path)
 	Repo.init(path)
 	r = Repo(path)
-	graphwalker = SimpleFetchGraphWalker(r.heads().values(), r.get_parents)
+	graphwalker = r.get_graph_walker()
 	f, commit = r.object_store.add_pack()
 	client.fetch_pack(host_path, r.object_store.determine_wants_all, 
 			          graphwalker, f.write, sys.stdout.write)

dulwich/client.py

@@ -44,38 +44,6 @@ def _fileno_can_read(fileno):
     return len(select.select([fileno], [], [], 0)[0]) > 0
 
 
-class SimpleFetchGraphWalker(object):
-    """Graph walker that finds out what commits are missing."""
-
-    def __init__(self, local_heads, get_parents):
-        """Create a new SimpleFetchGraphWalker instance.
-
-        :param local_heads: SHA1s that should be retrieved
-        :param get_parents: Function for finding the parents of a SHA1.
-        """
-        self.heads = set(local_heads)
-        self.get_parents = get_parents
-        self.parents = {}
-
-    def ack(self, sha):
-        """Ack that a particular revision and its ancestors are present in the target."""
-        if sha in self.heads:
-            self.heads.remove(sha)
-        if sha in self.parents:
-            for p in self.parents[sha]:
-                self.ack(p)
-
-    def next(self):
-        """Iterate over revisions that might be missing in the target."""
-        if self.heads:
-            ret = self.heads.pop()
-            ps = self.get_parents(ret)
-            self.parents[ret] = ps
-            self.heads.update(ps)
-            return ret
-        return None
-
-
 CAPABILITIES = ["multi_ack", "side-band-64k", "ofs-delta"]
 
 

dulwich/object_store.py

@@ -111,6 +111,22 @@ class BaseObjectStore(object):
         """
         return iter(MissingObjectFinder(self, wants, graph_walker, progress).next, None)
 
+    def get_commit_parents(self, sha):
+        """Retrieve the parents of a commit.
+
+        :param sha: SHA1 of the commit
+        :return: List of parent SHA1s
+        """
+        return self[sha].parents
+
+    def get_graph_walker(self, heads):
+        """Obtain a graph walker for this object store.
+        
+        :param heads: Local heads to start search with
+        :return: GraphWalker object
+        """
+        return ObjectStoreGraphWalker(heads, self.get_commit_parents)
+
 
 class DiskObjectStore(BaseObjectStore):
     """Git-style object store that exists on disk."""
@@ -476,3 +492,38 @@ class MissingObjectFinder(object):
         self.sha_done.add(sha)
         self.progress("counting objects: %d\r" % len(self.sha_done))
         return (sha, name)
+
+
+class ObjectStoreGraphWalker(object):
+    """Graph walker that finds out what commits are missing from an object store."""
+
+    def __init__(self, local_heads, get_parents):
+        """Create a new instance.
+
+        :param local_heads: Heads to start search with
+        :param get_parents: Function for finding the parents of a SHA1.
+        """
+        self.heads = set(local_heads)
+        self.get_parents = get_parents
+        self.parents = {}
+
+    def ack(self, sha):
+        """Ack that a particular revision and its ancestors are present in the source."""
+        if sha in self.heads:
+            self.heads.remove(sha)
+        if sha in self.parents:
+            for p in self.parents[sha]:
+                self.ack(p)
+
+    def next(self):
+        """Iterate over ancestors of heads in the target."""
+        if self.heads:
+            ret = self.heads.pop()
+            ps = self.get_parents(ret)
+            self.parents[ret] = ps
+            self.heads.update(ps)
+            return ret
+        return None
+
+
+

dulwich/repo.py

@@ -156,6 +156,11 @@ class Repo(object):
         return self.object_store.iter_shas(
             self.find_missing_objects(determine_wants, graph_walker, progress))
 
+    def get_graph_walker(self, heads=None):
+        if heads is None:
+            heads = self.heads().values()
+        return self.object_store.get_graph_walker(heads)
+
     def object_dir(self):
         """Return path of the object directory."""
         return os.path.join(self.controldir(), OBJECTDIR)
@@ -292,7 +297,7 @@ class Repo(object):
         return self.object_store[sha]
 
     def get_parents(self, sha):
-        return self.commit(sha).parents
+        return self.object_store.get_commit_parents(sha)
 
     def commit(self, sha):
         return self._get_object(sha, Commit)