diff --git a/src/anytree/iterators/preorderiter.py b/src/anytree/iterators/preorderiter.py index b07cf38..9e0d2e1 100644 --- a/src/anytree/iterators/preorderiter.py +++ b/src/anytree/iterators/preorderiter.py @@ -40,11 +40,18 @@ class PreOrderIter(AbstractIter): @staticmethod def _iter(children, filter_, stop, maxlevel): - for child_ in children: + cur_children = list(reversed(children)) + descendantmaxlevel = maxlevel + while cur_children: + child_ = cur_children.pop() if stop(child_): continue if filter_(child_): yield child_ - if not AbstractIter._abort_at_level(2, maxlevel): - descendantmaxlevel = maxlevel - 1 if maxlevel else None + if not AbstractIter._abort_at_level(2, descendantmaxlevel): + descendantmaxlevel = descendantmaxlevel - 1 if descendantmaxlevel else None + if len(child_.children) == 1: + cur_children.append(child_.children[0]) + continue yield from PreOrderIter._iter(child_.children, filter_, stop, descendantmaxlevel) + descendantmaxlevel = maxlevel diff --git a/tests/test_iterators.py b/tests/test_iterators.py index 78c8f1e..22fb0be 100644 --- a/tests/test_iterators.py +++ b/tests/test_iterators.py @@ -22,17 +22,23 @@ def test_preorder(): g = Node("g", parent=f) i = Node("i", parent=g) h = Node("h", parent=i) + k = Node("k", parent=a) + o = Node("o", parent=k) + m = Node("m", parent=e) - eq_(list(PreOrderIter(f)), [f, b, a, d, c, e, g, i, h]) + eq_(list(PreOrderIter(f)), [f, b, a, k, o, d, c, e, m, g, i, h]) eq_(list(PreOrderIter(f, maxlevel=0)), []) + eq_(list(PreOrderIter(f, maxlevel=2)), [f, b, g]) eq_(list(PreOrderIter(f, maxlevel=3)), [f, b, a, d, g, i]) - eq_(list(PreOrderIter(f, filter_=lambda n: n.name not in ("e", "g"))), [f, b, a, d, c, i, h]) - eq_(list(PreOrderIter(f, stop=lambda n: n.name == "d")), [f, b, a, g, i, h]) + eq_(list(PreOrderIter(f, maxlevel=4)), [f, b, a, k, d, c, e, g, i, h]) + eq_(list(PreOrderIter(f, maxlevel=f.height + 1)), list(PreOrderIter(f))) + eq_(list(PreOrderIter(f, filter_=lambda n: n.name not in ("e", "g"))), [f, b, a, k, o, d, c, m, i, h]) + eq_(list(PreOrderIter(f, stop=lambda n: n.name == "d")), [f, b, a, k, o, g, i, h]) it = PreOrderIter(f) eq_(next(it), f) eq_(next(it), b) - eq_(list(it), [a, d, c, e, g, i, h]) + eq_(list(it), [a, k, o, d, c, e, m, g, i, h]) def test_postorder():