From 3da0dbc5e84748fdd7db5ee78b9cd40636f19e7e Mon Sep 17 00:00:00 2001 From: Lukas Einhaus Date: Thu, 14 Nov 2024 11:28:34 +0100 Subject: [PATCH] fix(ir): fix bug where iterator was not remembering visited nodes --- elasticai/creator/ir/graph_delegate_test.py | 5 +---- elasticai/creator/ir/graph_iterators.py | 1 + 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/elasticai/creator/ir/graph_delegate_test.py b/elasticai/creator/ir/graph_delegate_test.py index 4541bfe2..f7c6e822 100644 --- a/elasticai/creator/ir/graph_delegate_test.py +++ b/elasticai/creator/ir/graph_delegate_test.py @@ -30,10 +30,7 @@ def test_iterating_breadth_first_upwards(): ) actual = tuple(bfs_iter_up(g.get_predecessors, "5")) - assert set(actual[0:2]) == {"3", "6"} - assert (set(actual[2:4]) == {"1", "2"} and actual[4] == "4") or ( - set(actual[3:5]) == {"1", "2"} and actual[2] == "4" - ) + assert actual == ("3", "6", "1", "2", "4", "0") def test_iterating_depth_first_preorder(): diff --git a/elasticai/creator/ir/graph_iterators.py b/elasticai/creator/ir/graph_iterators.py index f65f0937..c405462a 100644 --- a/elasticai/creator/ir/graph_iterators.py +++ b/elasticai/creator/ir/graph_iterators.py @@ -27,6 +27,7 @@ def bfs_iter_down(successors: NodeNeighbourFn, start: HashableT) -> Iterator[Has for p in successors(current): if p not in visited: yield p + visited.add(p) visit_next.append(p) visited.add(current)