From c0dcad78dd36796422a5534aaa19370b96f6893e Mon Sep 17 00:00:00 2001
From: Dave Borowitz <dborowitz@google.com>
Date: Fri, 30 Jul 2010 13:09:12 +0200
Subject: [PATCH] object_store: Include subtrees in iteration.

---
 NEWS                               |  3 +++
 dulwich/object_store.py            | 12 +++++++-----
 dulwich/tests/test_object_store.py | 28 ++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/NEWS b/NEWS
index 1955d9c..89bb03d 100644
--- a/NEWS
+++ b/NEWS
@@ -22,6 +22,9 @@
   * ObjectStore.iter_tree_contents now walks contents in depth-first, sorted
     order. (Dave Borowitz)
 
+  * ObjectStore.iter_tree_contents can optionally yield tree objects as well.
+    (Dave Borowitz).
+
 
 0.6.1	2010-07-22
 
diff --git a/dulwich/object_store.py b/dulwich/object_store.py
index 1a83439..725504f 100644
--- a/dulwich/object_store.py
+++ b/dulwich/object_store.py
@@ -175,24 +175,26 @@ class BaseObjectStore(object):
                     else:
                         todo.add((None, newhexsha, childpath))
 
-    def iter_tree_contents(self, tree_id):
+    def iter_tree_contents(self, tree_id, include_trees=False):
         """Iterate the contents of a tree and all subtrees.
 
-        Iteration is depth-first, as in e.g. os.walk.
+        Iteration is depth-first pre-order, as in e.g. os.walk.
 
         :param tree_id: SHA1 of the tree.
+        :param include_trees: If True, include tree objects in the iteration.
         :yield: Tuples of (path, mode, hexhsa) for objects in a tree.
         """
         todo = [('', stat.S_IFDIR, tree_id)]
         while todo:
             path, mode, hexsha = todo.pop()
-            if stat.S_ISDIR(mode):
+            is_subtree = stat.S_ISDIR(mode)
+            if not is_subtree or include_trees:
+                yield path, mode, hexsha
+            if is_subtree:
                 entries = reversed(self[hexsha].iteritems())
                 for name, entry_mode, entry_hexsha in entries:
                     entry_path = posixpath.join(path, name)
                     todo.append((entry_path, entry_mode, entry_hexsha))
-            else:
-                yield path, mode, hexsha
 
     def find_missing_objects(self, haves, wants, progress=None,
                              get_tagged=None):
diff --git a/dulwich/tests/test_object_store.py b/dulwich/tests/test_object_store.py
index e53a612..e34e75f 100644
--- a/dulwich/tests/test_object_store.py
+++ b/dulwich/tests/test_object_store.py
@@ -96,6 +96,34 @@ class ObjectStoreTests(object):
         self.assertEquals([(p, m, h) for (p, h, m) in blobs],
                           list(self.store.iter_tree_contents(tree_id)))
 
+    def test_iter_tree_contents_include_trees(self):
+        blob_a = make_object(Blob, data='a')
+        blob_b = make_object(Blob, data='b')
+        blob_c = make_object(Blob, data='c')
+        for blob in [blob_a, blob_b, blob_c]:
+            self.store.add_object(blob)
+
+        blobs = [
+          ('a', blob_a.id, 0100644),
+          ('ad/b', blob_b.id, 0100644),
+          ('ad/bd/c', blob_c.id, 0100755),
+          ]
+        tree_id = commit_tree(self.store, blobs)
+        tree = self.store[tree_id]
+        tree_ad = self.store[tree['ad'][1]]
+        tree_bd = self.store[tree_ad['bd'][1]]
+
+        expected = [
+          ('', 0040000, tree_id),
+          ('a', 0100644, blob_a.id),
+          ('ad', 0040000, tree_ad.id),
+          ('ad/b', 0100644, blob_b.id),
+          ('ad/bd', 0040000, tree_bd.id),
+          ('ad/bd/c', 0100755, blob_c.id),
+          ]
+        actual = self.store.iter_tree_contents(tree_id, include_trees=True)
+        self.assertEquals(expected, list(actual))
+
 
 class MemoryObjectStoreTests(ObjectStoreTests, TestCase):
 
-- 
2.34.1