Improve standard visitor support

Author: qballand

freexgraph/freexgraph.py

@@ -39,15 +39,20 @@ class FreExNode:
     parents: Set[str]
     """Parents of the node to add """
 
+    extension_node: bool
+    """ """
+
     def __init__(
         self,
         uid: str = None,
         *,
         fork_id: Optional[str] = None,
         parents: Set[str] = None,
         graph_ref: nx.DiGraph = None,
+        extension_node: bool = False,
     ):
         self.parents = parents or set()
+        self.extension_node = extension_node
         self._graph_ref = graph_ref
         self._id = uid
         self._fork_id = fork_id
@@ -63,6 +68,8 @@ def apply_accept_(self, visitor: AnyVisitor) -> bool:
         from freexgraph.standard_visitor import is_standard_visitor
 
         if is_standard_visitor(visitor):
+            if self.extension_node:
+                self.accept(visitor)
             return visitor.visit_standard(self)
         return self.accept(visitor)
 
@@ -126,7 +133,9 @@ class GraphNode(FreExNode):
     def __init__(
         self, uid: str = None, *, graph: "FreExGraph", parents: Set[str] = None
     ):
-        super().__init__(uid=uid, parents=parents, graph_ref=graph._graph)
+        super().__init__(
+            uid=uid, parents=parents, graph_ref=graph._graph, extension_node=True
+        )
         self._graph_ex = graph
 
     def accept(self, visitor: AnyVisitor) -> bool:

test/basic_freexgraph_test.py

@@ -207,6 +207,12 @@ def test_graph_node(valid_basic_execution_graph, visitor_test):
     assert visitor_test.inner_graph_started[0].startswith(id_graph)
     assert visitor_test.inner_graph_started == visitor_test.inner_graph_ended
 
+    # test find in graph node from root
+    finder = FindFirstVisitor(lambda n: n.id.startswith("id3"))
+    finder.visit(execution_graph.root)
+    assert finder.found()
+    assert finder.result.id.startswith("id3")
+
 
 def test_add_nodes(node_list_complex_graph, visitor_test):
     execution_graph = FreExGraph()