Add: iterator over leaf keys in a CowTree

2 files changed, 16 insertions, 0 deletions

diff --git a/obnamlib/fmt_ga/cowtree.py b/obnamlib/fmt_ga/cowtree.py
index f1f74385..69b00a9e 100644
--- a/obnamlib/fmt_ga/cowtree.py
+++ b/obnamlib/fmt_ga/cowtree.py
@@ -85,6 +85,12 @@ class CowTree(object):
         new_id = self._store.put_leaf(new)
         self._leaf_list.insert_leaf(sorted_keys[0], sorted_keys[-1], new_id)
 
+    def keys(self):
+        for leaf_id in self._leaf_list.leaf_ids():
+            leaf = self._store.get_leaf(leaf_id)
+            for key in leaf.keys():
+                yield key
+
     def commit(self):
         fake_leaf = obnamlib.CowLeaf()
         fake_leaf.insert('leaf_list', self._leaf_list.as_dict())
@@ -106,6 +112,10 @@ class _LeafList(object):
     def from_dict(self, some_dict):
         self._leaf_list = some_dict
 
+    def leaf_ids(self):
+        for leaf_info in self._leaf_list:
+            yield leaf_info['id']
+
     def find_leaf_for_key(self, key):
         # If there are no leaves, we can't pick one for key.
         if not self._leaf_list:
diff --git a/obnamlib/fmt_ga/cowtree_tests.py b/obnamlib/fmt_ga/cowtree_tests.py
index b34ce2f0..c53a9ab7 100644
--- a/obnamlib/fmt_ga/cowtree_tests.py
+++ b/obnamlib/fmt_ga/cowtree_tests.py
@@ -64,3 +64,9 @@ class CowTreeTests(unittest.TestCase):
         cow2.set_leaf_store(self.ls)
         cow2.set_list_node(list_id)
         self.assertEqual(cow2.lookup(key), value)
+
+    def test_iterates_over_leaf_keys(self):
+        key = 'fookey'
+        value = 'barvalue'
+        self.cow.insert(key, value)
+        self.assertEqual(list(self.cow.keys()), [key])