# -*- coding: utf-8 -*-
"""
Tree Iteration.
* :any:`PreOrderIter`: iterate over tree using pre-order strategy (self, children)
* :any:`PostOrderIter`: iterate over tree using post-order strategy (children, self)
* :any:`LevelOrderIter`: iterate over tree using level-order strategy
* :any:`LevelOrderGroupIter`: iterate over tree using level-order strategy returning group for every level
* :any:`ZigZagGroupIter`: iterate over tree using level-order strategy returning group for every level
"""
import six
[docs]class AbstractIter(six.Iterator):
def __init__(self, node, filter_=None, stop=None, maxlevel=None):
"""
Base class for all iterators.
Iterate over tree starting at `node`.
Keyword Args:
filter_: function called with every `node` as argument, `node` is returned if `True`.
stop: stop iteration at `node` if `stop` function returns `True` for `node`.
maxlevel (int): maximum decending in the node hierarchy.
"""
self.node = node
self.filter_ = filter_
self.stop = stop
self.maxlevel = maxlevel
self.__iter = None
def __init(self):
node = self.node
filter_ = self.filter_
stop = self.stop
maxlevel = self.maxlevel
if filter_ is None:
def filter_(node):
return True
if stop is None:
def stop(node):
return False
children = [] if AbstractIter._abort_at_level(1, maxlevel) else AbstractIter._get_children([node], stop)
return self._iter(children, filter_, stop, maxlevel)
def __iter__(self):
return self.__init()
def __next__(self):
if self.__iter is None:
self.__iter = self.__init()
return next(self.__iter)
@staticmethod
def _iter(children, filter_, stop, maxlevel):
raise NotImplementedError() # pragma: no cover
@staticmethod
def _abort_at_level(level, maxlevel):
return maxlevel is not None and level > maxlevel
@staticmethod
def _get_children(children, stop):
return [child for child in children if not stop(child)]
[docs]class PreOrderIter(AbstractIter):
"""
Iterate over tree applying pre-order strategy starting at `node`.
Start at root and go-down until reaching a leaf node.
Step upwards then, and search for the next leafs.
>>> from anytree import Node, RenderTree, AsciiStyle
>>> f = Node("f")
>>> b = Node("b", parent=f)
>>> a = Node("a", parent=b)
>>> d = Node("d", parent=b)
>>> c = Node("c", parent=d)
>>> e = Node("e", parent=d)
>>> g = Node("g", parent=f)
>>> i = Node("i", parent=g)
>>> h = Node("h", parent=i)
>>> print(RenderTree(f, style=AsciiStyle()).by_attr())
f
|-- b
| |-- a
| +-- d
| |-- c
| +-- e
+-- g
+-- i
+-- h
>>> [node.name for node in PreOrderIter(f)]
['f', 'b', 'a', 'd', 'c', 'e', 'g', 'i', 'h']
>>> [node.name for node in PreOrderIter(f, maxlevel=3)]
['f', 'b', 'a', 'd', 'g', 'i']
>>> [node.name for node in PreOrderIter(f, filter_=lambda n: n.name not in ('e', 'g'))]
['f', 'b', 'a', 'd', 'c', 'i', 'h']
>>> [node.name for node in PreOrderIter(f, stop=lambda n: n.name == 'd')]
['f', 'b', 'a', 'g', 'i', 'h']
"""
@staticmethod
def _iter(children, filter_, stop, maxlevel):
stack = [children]
while stack:
children = stack[-1]
if children:
child = children.pop(0)
if filter_(child):
yield child
if not AbstractIter._abort_at_level(len(stack) + 1, maxlevel):
grandchildren = AbstractIter._get_children(child.children, stop)
if grandchildren:
stack.append(grandchildren)
else:
stack.pop()
[docs]class PostOrderIter(AbstractIter):
"""
Iterate over tree applying post-order strategy starting at `node`.
>>> from anytree import Node, RenderTree, AsciiStyle
>>> f = Node("f")
>>> b = Node("b", parent=f)
>>> a = Node("a", parent=b)
>>> d = Node("d", parent=b)
>>> c = Node("c", parent=d)
>>> e = Node("e", parent=d)
>>> g = Node("g", parent=f)
>>> i = Node("i", parent=g)
>>> h = Node("h", parent=i)
>>> print(RenderTree(f, style=AsciiStyle()).by_attr())
f
|-- b
| |-- a
| +-- d
| |-- c
| +-- e
+-- g
+-- i
+-- h
>>> [node.name for node in PostOrderIter(f)]
['a', 'c', 'e', 'd', 'b', 'h', 'i', 'g', 'f']
>>> [node.name for node in PostOrderIter(f, maxlevel=3)]
['a', 'd', 'b', 'i', 'g', 'f']
>>> [node.name for node in PostOrderIter(f, filter_=lambda n: n.name not in ('e', 'g'))]
['a', 'c', 'd', 'b', 'h', 'i', 'f']
>>> [node.name for node in PostOrderIter(f, stop=lambda n: n.name == 'd')]
['a', 'b', 'h', 'i', 'g', 'f']
"""
@staticmethod
def _iter(children, filter_, stop, maxlevel):
return PostOrderIter.__next(children, 1, filter_, stop, maxlevel)
@staticmethod
def __next(children, level, filter_, stop, maxlevel):
if not AbstractIter._abort_at_level(level, maxlevel):
for child in children:
grandchildren = AbstractIter._get_children(child.children, stop)
for grandchild in PostOrderIter.__next(grandchildren, level + 1, filter_, stop, maxlevel):
yield grandchild
if filter_(child):
yield child
[docs]class LevelOrderIter(AbstractIter):
"""
Iterate over tree applying level-order strategy starting at `node`.
>>> from anytree import Node, RenderTree, AsciiStyle
>>> f = Node("f")
>>> b = Node("b", parent=f)
>>> a = Node("a", parent=b)
>>> d = Node("d", parent=b)
>>> c = Node("c", parent=d)
>>> e = Node("e", parent=d)
>>> g = Node("g", parent=f)
>>> i = Node("i", parent=g)
>>> h = Node("h", parent=i)
>>> print(RenderTree(f, style=AsciiStyle()).by_attr())
f
|-- b
| |-- a
| +-- d
| |-- c
| +-- e
+-- g
+-- i
+-- h
>>> [node.name for node in LevelOrderIter(f)]
['f', 'b', 'g', 'a', 'd', 'i', 'c', 'e', 'h']
>>> [node.name for node in LevelOrderIter(f, maxlevel=3)]
['f', 'b', 'g', 'a', 'd', 'i']
>>> [node.name for node in LevelOrderIter(f, filter_=lambda n: n.name not in ('e', 'g'))]
['f', 'b', 'a', 'd', 'i', 'c', 'h']
>>> [node.name for node in LevelOrderIter(f, stop=lambda n: n.name == 'd')]
['f', 'b', 'g', 'a', 'i', 'h']
"""
@staticmethod
def _iter(children, filter_, stop, maxlevel):
level = 1
while children:
next_children = []
for child in children:
if filter_(child):
yield child
next_children += AbstractIter._get_children(child.children, stop)
children = next_children
level += 1
if AbstractIter._abort_at_level(level, maxlevel):
break
[docs]class LevelOrderGroupIter(AbstractIter):
"""
Iterate over tree applying level-order strategy with grouping starting at `node`.
Return a tuple of nodes for each level. The first tuple contains the
nodes at level 0 (always `node`). The second tuple contains the nodes at level 1
(children of `node`). The next level contains the children of the children, and so on.
>>> from anytree import Node, RenderTree, AsciiStyle
>>> f = Node("f")
>>> b = Node("b", parent=f)
>>> a = Node("a", parent=b)
>>> d = Node("d", parent=b)
>>> c = Node("c", parent=d)
>>> e = Node("e", parent=d)
>>> g = Node("g", parent=f)
>>> i = Node("i", parent=g)
>>> h = Node("h", parent=i)
>>> print(RenderTree(f, style=AsciiStyle()).by_attr())
f
|-- b
| |-- a
| +-- d
| |-- c
| +-- e
+-- g
+-- i
+-- h
>>> [[node.name for node in children] for children in LevelOrderGroupIter(f)]
[['f'], ['b', 'g'], ['a', 'd', 'i'], ['c', 'e', 'h']]
>>> [[node.name for node in children] for children in LevelOrderGroupIter(f, maxlevel=3)]
[['f'], ['b', 'g'], ['a', 'd', 'i']]
>>> [[node.name for node in children]
... for children in LevelOrderGroupIter(f, filter_=lambda n: n.name not in ('e', 'g'))]
[['f'], ['b'], ['a', 'd', 'i'], ['c', 'h']]
>>> [[node.name for node in children]
... for children in LevelOrderGroupIter(f, stop=lambda n: n.name == 'd')]
[['f'], ['b', 'g'], ['a', 'i'], ['h']]
"""
@staticmethod
def _iter(children, filter_, stop, maxlevel):
level = 1
while children:
yield tuple([child for child in children if filter_(child)])
level += 1
if AbstractIter._abort_at_level(level, maxlevel):
break
children = LevelOrderGroupIter._get_grandchildren(children, stop)
@staticmethod
def _get_grandchildren(children, stop):
next_children = []
for child in children:
next_children = next_children + AbstractIter._get_children(child.children, stop)
return next_children
[docs]class ZigZagGroupIter(AbstractIter):
"""
Iterate over tree applying Zig-Zag strategy with grouping starting at `node`.
Return a tuple of nodes for each level. The first tuple contains the
nodes at level 0 (always `node`). The second tuple contains the nodes at level 1
(children of `node`) in reversed order.
The next level contains the children of the children in forward order, and so on.
>>> from anytree import Node, RenderTree, AsciiStyle
>>> f = Node("f")
>>> b = Node("b", parent=f)
>>> a = Node("a", parent=b)
>>> d = Node("d", parent=b)
>>> c = Node("c", parent=d)
>>> e = Node("e", parent=d)
>>> g = Node("g", parent=f)
>>> i = Node("i", parent=g)
>>> h = Node("h", parent=i)
>>> print(RenderTree(f, style=AsciiStyle()).by_attr())
f
|-- b
| |-- a
| +-- d
| |-- c
| +-- e
+-- g
+-- i
+-- h
>>> [[node.name for node in children] for children in ZigZagGroupIter(f)]
[['f'], ['g', 'b'], ['a', 'd', 'i'], ['h', 'e', 'c']]
>>> [[node.name for node in children] for children in ZigZagGroupIter(f, maxlevel=3)]
[['f'], ['g', 'b'], ['a', 'd', 'i']]
>>> [[node.name for node in children]
... for children in ZigZagGroupIter(f, filter_=lambda n: n.name not in ('e', 'g'))]
[['f'], ['b'], ['a', 'd', 'i'], ['h', 'c']]
>>> [[node.name for node in children]
... for children in ZigZagGroupIter(f, stop=lambda n: n.name == 'd')]
[['f'], ['g', 'b'], ['a', 'i'], ['h']]
"""
@staticmethod
def _iter(children, filter_, stop, maxlevel):
_iter = LevelOrderGroupIter._iter(children, filter_, stop, maxlevel)
while True:
yield next(_iter)
yield tuple(reversed(next(_iter)))