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
This commit is contained in:
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user