Skip to content
Snippets Groups Projects
Commit 153c6e84 authored by Daniele Serafini's avatar Daniele Serafini
Browse files

make deduplication optional when iterating over the merkle tree

parent cfb3073c
No related branches found
Tags v2.6.4
No related merge requests found
......@@ -277,18 +277,20 @@ class MerkleNode(dict, metaclass=abc.ABCMeta):
for child in self.values():
child.reset_collect()
def iter_tree(self) -> Iterator["MerkleNode"]:
"""Yields all children nodes, recursively. Common nodes are
deduplicated.
def iter_tree(self, dedup=True) -> Iterator["MerkleNode"]:
"""Yields all children nodes, recursively. Common nodes are deduplicated
by default (deduplication can be turned off setting the given argument
'dedup' to False).
"""
yield from self._iter_tree(set())
yield from self._iter_tree(set(), dedup)
def _iter_tree(self, seen: Set[bytes]) -> Iterator["MerkleNode"]:
def _iter_tree(self, seen: Set[bytes], dedup) -> Iterator["MerkleNode"]:
if self.hash not in seen:
seen.add(self.hash)
if dedup:
seen.add(self.hash)
yield self
for child in self.values():
yield from child._iter_tree(seen=seen)
yield from child._iter_tree(seen=seen, dedup=dedup)
class MerkleLeaf(MerkleNode):
......
......@@ -172,10 +172,18 @@ class TestMerkleNode(unittest.TestCase):
collected2 = self.root.collect()
self.assertEqual(collected2, {})
def test_iter_tree(self):
def test_iter_tree_with_deduplication(self):
nodes = list(self.root.iter_tree())
self.assertCountEqual(nodes, self.nodes.values())
def test_iter_tree_without_deduplication(self):
# duplicate existing hash in merkle tree
self.root[b"d"] = MerkleTestNode({"value": b"root/c/c/c"})
nodes_dedup = list(self.root.iter_tree())
nodes = list(self.root.iter_tree(dedup=False))
assert nodes != nodes_dedup
assert len(nodes) == len(nodes_dedup) + 1
def test_get(self):
for key in (b"a", b"b", b"c"):
self.assertEqual(self.root[key], self.nodes[b"root/" + key])
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment