From 19e0789ebe76647b63f9fce976edff70a15e8103 Mon Sep 17 00:00:00 2001 From: Joshua Harlow Date: Tue, 10 Feb 2015 15:36:03 -0800 Subject: [PATCH] Add a BFS tree iterator To complement the DFS iterator that we provide in our tree helper/type class add a BFS iterator as well so that when/if we extract this module to elsewhere we can have a nice featureful (and hopefully useful) set of iterators for folks to use. Change-Id: I1bc10c25bb62a5ffc85863f3f894d5469db95ff6 --- taskflow/tests/unit/test_types.py | 9 +++++++++ taskflow/types/tree.py | 26 ++++++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/taskflow/tests/unit/test_types.py b/taskflow/tests/unit/test_types.py index 83d938d5..2151b9e1 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 2386bc9c..f07ff9a0 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.""" @@ -198,3 +220,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)