diff options
author | Joshua Harlow <harlowja@yahoo-inc.com> | 2015-02-10 15:36:03 -0800 |
---|---|---|
committer | Joshua Harlow <harlowja@yahoo-inc.com> | 2015-02-10 15:36:03 -0800 |
commit | 19e0789ebe76647b63f9fce976edff70a15e8103 (patch) | |
tree | 62b4a3e61a5e71ada586d23090fdf7c8ebfd49c4 | |
parent | 3bf3249cc7103dcdcbdc3881a546a6cad8e4b65d (diff) | |
download | taskflow-19e0789ebe76647b63f9fce976edff70a15e8103.tar.gz |
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
-rw-r--r-- | taskflow/tests/unit/test_types.py | 9 | ||||
-rw-r--r-- | taskflow/types/tree.py | 26 |
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 2386bc9..f07ff9a 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) |