Allow node finding to not do a deep search

When a local search is preferred allow for it
to be done by passing in a 'only_direct' option,
also allows passing in a 'include_self' option
to ignore/include the current node when finding.

The defaults of 'include_self' and 'only_direct'
work just as they did prior to this change, so that
the addition does not affect current usage(s) of
this method.

Change-Id: I89fb476bee32144092ea141ca3b971118ab633be
This commit is contained in:
Joshua Harlow
2015-03-03 12:17:01 -08:00
parent d4f02f6fea
commit 53af5d3221
2 changed files with 23 additions and 2 deletions

View File

@@ -136,6 +136,16 @@ class TreeTest(test.TestCase):
root.freeze()
self.assertRaises(tree.FrozenNode, root.add, "bird")
def test_find(self):
root = self._make_species()
self.assertIsNone(root.find('monkey', only_direct=True))
self.assertIsNotNone(root.find('monkey', only_direct=False))
self.assertIsNotNone(root.find('animal', only_direct=True))
self.assertIsNotNone(root.find('reptile', only_direct=True))
self.assertIsNone(root.find('animal', include_self=False))
self.assertIsNone(root.find('animal',
include_self=False, only_direct=True))
def test_dfs_itr(self):
root = self._make_species()
things = list([n.item for n in root.dfs_iter(include_self=True)])

View File

@@ -17,6 +17,7 @@
# under the License.
import collections
import itertools
import os
import six
@@ -118,7 +119,7 @@ class Node(object):
yield node
node = node.parent
def find(self, item):
def find(self, item, only_direct=False, include_self=True):
"""Returns the node for an item if it exists in this node.
This will search not only this node but also any children nodes and
@@ -126,9 +127,19 @@ class Node(object):
object.
:param item: item to lookup.
:param only_direct: only look at current node and its direct children.
:param include_self: include the current node during searching.
:returns: the node for an item if it exists in this node
"""
for n in self.dfs_iter(include_self=True):
if only_direct:
if include_self:
it = itertools.chain([self], self.reverse_iter())
else:
it = self.reverse_iter()
else:
it = self.dfs_iter(include_self=include_self)
for n in it:
if n.item == item:
return n
return None