Merge "Add a BFS tree iterator"

2 files changed, 35 insertions, 0 deletions

diff --git a/taskflow/tests/unit/test_types.py b/taskflow/tests/unit/test_types.py
index 83d938d..2151b9e 100644
--- a/taskflow/tests/unit/test_types.py
+++ b/taskflow/tests/unit/test_types.py
@@ -151,6 +151,15 @@ class TreeTest(test.TestCase):
         self.assertEqual(['mammal', 'horse', 'primate',
                           'monkey', 'human', 'reptile'], things)
 
+    def test_bfs_iter(self):
+        root = self._make_species()
+        things = list([n.item for n in root.bfs_iter(include_self=True)])
+        self.assertEqual(['animal', 'reptile', 'mammal', 'primate',
+                          'horse', 'human', 'monkey'], things)
+        things = list([n.item for n in root.bfs_iter(include_self=False)])
+        self.assertEqual(['reptile', 'mammal', 'primate',
+                          'horse', 'human', 'monkey'], things)
+
 
 class StopWatchTest(test.TestCase):
     def setUp(self):
diff --git a/taskflow/types/tree.py b/taskflow/types/tree.py
index 30edd45..c427314 100644
--- a/taskflow/types/tree.py
+++ b/taskflow/types/tree.py
@@ -16,6 +16,7 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
+import collections
 import os
 
 import six
@@ -49,6 +50,27 @@ class _DFSIter(object):
             stack.extend(node.reverse_iter())
 
 
+class _BFSIter(object):
+    """Breadth first iterator (non-recursive) over the child nodes."""
+
+    def __init__(self, root, include_self=False):
+        self.root = root
+        self.include_self = bool(include_self)
+
+    def __iter__(self):
+        q = collections.deque()
+        if self.include_self:
+            q.append(self.root)
+        else:
+            q.extend(self.root.reverse_iter())
+        while q:
+            node = q.popleft()
+            # Visit the node.
+            yield node
+            # Traverse the left & right subtree.
+            q.extend(node.reverse_iter())
+
+
 class Node(object):
     """A n-ary node class that can be used to create tree structures."""
 
@@ -206,3 +228,7 @@ class Node(object):
     def dfs_iter(self, include_self=False):
         """Depth first iteration (non-recursive) over the child nodes."""
         return _DFSIter(self, include_self=include_self)
+
+    def bfs_iter(self, include_self=False):
+        """Breadth first iteration (non-recursive) over the child nodes."""
+        return _BFSIter(self, include_self=include_self)