From 9cf7a04a3e0c3d64fdc79419f631cd8fd0461c3c Mon Sep 17 00:00:00 2001
From: Valentin Lorentz <vlorentz@softwareheritage.org>
Date: Mon, 24 Feb 2020 15:59:14 +0100
Subject: [PATCH] Add method MerkleNode.iter_tree, to visit all nodes in the
subtree of a node.
---
swh/model/merkle.py | 16 +++++++++++++++-
swh/model/tests/test_merkle.py | 4 ++++
2 files changed, 19 insertions(+), 1 deletion(-)
diff --git a/swh/model/merkle.py b/swh/model/merkle.py
index 31be3d17..9d97efdc 100644
--- a/swh/model/merkle.py
+++ b/swh/model/merkle.py
@@ -8,7 +8,7 @@
import abc
import collections
-from typing import List, Optional
+from typing import Iterator, List, Optional, Set
def deep_update(left, right):
@@ -273,6 +273,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.
+ """
+ yield from self._iter_tree(set())
+
+ def _iter_tree(
+ self, seen: Set[bytes]) -> Iterator['MerkleNode']:
+ if self.hash not in seen:
+ seen.add(self.hash)
+ yield self
+ for child in self.values():
+ yield from child._iter_tree(seen=seen)
+
class MerkleLeaf(MerkleNode):
"""A leaf to a Merkle tree.
diff --git a/swh/model/tests/test_merkle.py b/swh/model/tests/test_merkle.py
index dc7da63b..734f7c03 100644
--- a/swh/model/tests/test_merkle.py
+++ b/swh/model/tests/test_merkle.py
@@ -184,6 +184,10 @@ class TestMerkleNode(unittest.TestCase):
collected2 = self.root.collect()
self.assertEqual(collected2, {})
+ def test_iter_tree(self):
+ nodes = list(self.root.iter_tree())
+ self.assertCountEqual(nodes, self.nodes.values())
+
def test_get(self):
for key in (b'a', b'b', b'c'):
self.assertEqual(self.root[key], self.nodes[b'root/' + key])
--
GitLab