improve and perfect topological sort functions

Author: ArnoChenFx

NodeGraphQt/__init__.py

@@ -89,13 +89,17 @@ def __init__(self):
 from .pkg_info import __license__ as LICENSE
 
 # functions
-from .base.utils import setup_context_menu
+from .base.utils import setup_context_menu, \
+    topological_sort_by_up, topological_sort_by_down, \
+    update_node_down_stream, update_node_up_stream,\
+    update_nodes_by_up, update_nodes_by_down
 
 # widgets
 from .widgets.node_tree import NodeTreeWidget
 from .widgets.properties_bin import PropertiesBinWidget
 from .widgets.node_publish_widget import NodePublishWidget
 
+
 __version__ = VERSION
 __all__ = [
     'BackdropNode',
@@ -114,5 +118,11 @@ def __init__(self):
     'setup_context_menu',
     'NodePublishWidget',
     'SubGraph',
+    'topological_sort_by_up',
+    'topological_sort_by_down',
+    'update_node_up_stream',
+    'update_node_down_stream',
+    'update_nodes_by_up',
+    'update_nodes_by_down',
 ]
 

NodeGraphQt/base/graph.py

@@ -839,8 +839,9 @@ def set_node_space(self, node):
         Args:
             node (NodeGraphQt.SubGraph): node object.
         """
-        if node is self._current_node_space:
+        if node is self._current_node_space or not isinstance(node, SubGraph):
             return
+
         if self._current_node_space is not None:
             self._current_node_space.exit()
 

NodeGraphQt/base/node.py

@@ -20,6 +20,7 @@
                                     NodeIntEdit,
                                     NodeCheckBox,
                                     NodeFilePath)
+from .utils import update_node_down_stream
 
 
 class classproperty(object):
@@ -1060,35 +1061,21 @@ def on_input_disconnected(self, in_port, out_port):
         """
         return
 
-    def update_streams(self, *args):
+    def update_stream(self):
         """
-        Update all nodes joined by pipes to this.
+        Update node down stream.
         """
-        nodes = []
-        trash = []
-
-        for port, nodeList in self.connected_output_nodes().items():
-            nodes.extend(nodeList)
-
-        while nodes:
-            node = nodes.pop()
-            if node.disabled():
-                continue
-            if node not in trash:
-                trash.append(node)
-
-            for port, nodeList in node.connected_output_nodes().items():
-                nodes.extend(nodeList)
-
-            if not node.connected_output_nodes():
-                try:
-                    node.run()
-                except Exception as error:
-                    print("Error Update Streams: %s" % str(error))
+        update_node_down_stream(self)
 
     def run(self):
         """
-        Node evaluation logics.
+        Node evaluation logic.
+        """
+        return
+
+    def when_disabled(self):
+        """
+        Node evaluation logic when node has been disabled.
         """
         return
 

NodeGraphQt/base/utils.py

@@ -2,7 +2,6 @@
 from distutils.version import LooseVersion
 
 from .. import QtGui, QtCore
-from .node import SubGraph
 from ..constants import (PIPE_LAYOUT_CURVED,
                          PIPE_LAYOUT_STRAIGHT,
                          PIPE_LAYOUT_ANGLE)
@@ -250,9 +249,7 @@ def _fit_to_selection(graph):
 def _jump_in(graph):
     nodes = graph.selected_nodes()
     if nodes:
-        node = nodes[0]
-        if isinstance(node, SubGraph):
-            graph.set_node_space(node)
+        graph.set_node_space(nodes[0])
 
 
 def _jump_out(graph):
@@ -312,7 +309,7 @@ def _has_output_node(node):
     return False
 
 
-def _build_graph(start_nodes):
+def _build_down_stream_graph(start_nodes):
     graph = {}
     for node in start_nodes:
         output_nodes = get_output_nodes(node)
@@ -325,19 +322,26 @@ def _build_graph(start_nodes):
                     graph[n] = nodes
                     _output_nodes.extend(nodes)
             output_nodes = _output_nodes
-
     return graph
 
 
-def topological_sort(start_nodes=[], all_nodes=[]):
-    if not start_nodes:
-        start_nodes = [n for n in all_nodes if not _has_input_node(n)]
-    if not start_nodes:
-        return []
-    if not [n for n in start_nodes if _has_output_node(n)]:
-        return start_nodes
+def _build_up_stream_graph(start_nodes):
+    graph = {}
+    for node in start_nodes:
+        input_nodes = get_input_nodes(node)
+        graph[node] = input_nodes
+        while input_nodes:
+            _input_nodes = []
+            for n in input_nodes:
+                if n not in graph:
+                    nodes = get_input_nodes(n)
+                    graph[n] = nodes
+                    _input_nodes.extend(nodes)
+            input_nodes = _input_nodes
+    return graph
+
 
-    graph = _build_graph(start_nodes)
+def _sort_nodes(graph, start_nodes, reverse=True):
     if not graph:
         return []
 
@@ -357,6 +361,57 @@ def dfs(graph, start_node):
             visit[start_node] = True
             dfs(graph, start_node)
 
-    sorted_nodes.reverse()
+    if reverse:
+        sorted_nodes.reverse()
 
     return sorted_nodes
+
+
+def topological_sort_by_down(start_nodes=[], all_nodes=[]):
+    if not start_nodes:
+        start_nodes = [n for n in all_nodes if not _has_input_node(n)]
+    if not start_nodes:
+        return []
+    if not [n for n in start_nodes if _has_output_node(n)]:
+        return start_nodes
+
+    graph = _build_down_stream_graph(start_nodes)
+
+    return _sort_nodes(graph, start_nodes, True)
+
+
+def topological_sort_by_up(start_nodes=[], all_nodes=[]):
+    if not start_nodes:
+        start_nodes = [n for n in all_nodes if not _has_output_node(n)]
+    if not start_nodes:
+        return []
+    if not [n for n in start_nodes if _has_input_node(n)]:
+        return start_nodes
+
+    graph = _build_up_stream_graph(start_nodes)
+
+    return _sort_nodes(graph, start_nodes, False)
+
+
+def _update_nodes(nodes):
+    for node in nodes:
+        if node.disabled():
+            node.when_disabled()
+        else:
+            node.run()
+
+
+def update_node_down_stream(node):
+    _update_nodes(topological_sort_by_down(start_nodes=[node]))
+
+
+def update_node_up_stream(node):
+    _update_nodes(topological_sort_by_up(start_nodes=[node]))
+
+
+def update_nodes_by_down(nodes):
+    _update_nodes(topological_sort_by_down(all_nodes=nodes))
+
+
+def update_nodes_by_up(nodes):
+    _update_nodes(topological_sort_by_up(all_nodes=nodes))

example_auto_nodes.py

@@ -5,13 +5,12 @@
 from NodeGraphQt import NodeGraph, setup_context_menu
 from NodeGraphQt import QtWidgets, QtCore, PropertiesBinWidget, \
     NodeTreeWidget, BackdropNode, NodePublishWidget
-from NodeGraphQt.base.utils import topological_sort
 import os
 import sys
 import inspect
 import importlib
 from example_auto_nodes import AutoNode, ModuleNode, \
-    SubGraphNode, Publish, RootNode
+    SubGraphNode, Publish, RootNode, update_nodes
 
 
 def get_nodes_from_folder(folder_path):
@@ -52,7 +51,7 @@ def get_published_nodes_from_folder(folder_path):
 
 
 def cook_node(graph, node):
-    node.cook(forceCook=True)
+    node.update_stream(forceCook=True)
 
 
 def print_functions(graph, node):
@@ -87,11 +86,7 @@ def publish_node(graph, node):
 
 
 def cook_nodes(nodes):
-    nodes = topological_sort(all_nodes=nodes)
-    for node in nodes:
-        node.cook(stream=True)
-        if node.error():
-            break
+    update_nodes(nodes)
 
 
 if __name__ == '__main__':

example_auto_nodes/__init__.py

@@ -2,3 +2,4 @@
 from .node_base.module_node import ModuleNode
 from .node_base.subgraph_node import SubGraphNode, RootNode
 from .subgraph_nodes import Publish
+from .node_base.utils import update_node_down_stream, update_nodes

example_auto_nodes/networks/example_SubGraph.json

@@ -262,7 +262,7 @@
     "0x169cc6f5ec8":{
       "type_":"Inputs.Vector3InputNode",
       "icon":null,
-      "name":"Vector 5",
+      "name":"Vector 6",
       "color":[
         13,
         18,

example_auto_nodes/node_base/auto_node.py

@@ -1,6 +1,6 @@
 from NodeGraphQt import BaseNode, SubGraph, Port, QtCore
 from NodeGraphQt.constants import NODE_PROP
-from NodeGraphQt.base.utils import topological_sort, get_output_nodes
+from . utils import update_node_down_stream
 import traceback
 import hashlib
 import copy
@@ -72,12 +72,13 @@ def cookTime(self, time):
         self._cookTime = time
         self._update_tool_tip()
 
-    def update_streams(self):
-        nodes = topological_sort(start_nodes=get_output_nodes(self))
-        for node in nodes:
-            node.cook(stream=True)
-            if node.error():
-                break
+    def update_stream(self, forceCook=False):
+        if not forceCook:
+            if not self._autoCook or not self.needCook:
+                return
+            if self.graph is not None and not self.graph.auto_update:
+                return
+        update_node_down_stream(self)
 
     def cookNextNode(self):
         nodes = []
@@ -124,24 +125,14 @@ def when_disabled(self):
         for index, out_port in enumerate(self.output_ports()):
             self.set_property(out_port.name(), self.getInputData(min(index, num)))
 
-    def cook(self, forceCook=False, stream=False):
-        if self.disabled() and stream:
-            self.when_disabled()
-            return
-        if not forceCook:
-            if not self._autoCook or not self.needCook:
-                return
-            if self.graph is not None and not self.graph.auto_update:
-                return
-
+    def cook(self):
         _tmp = self._autoCook
         self._autoCook = False
 
         if self.error():
             self._close_error()
 
         _start_time = time.time()
-
         try:
             self.run()
         except:
@@ -156,15 +147,12 @@ def cook(self, forceCook=False, stream=False):
 
         self.cooked.emit()
 
-        if not stream:
-            self.update_streams()
-
     def run(self):
         pass
 
     def on_input_connected(self, to_port, from_port):
         if self.checkPortType(to_port, from_port):
-            self.cook()
+            self.update_stream()
         else:
             self.needCook = False
             to_port.disconnect_from(from_port)
@@ -173,7 +161,7 @@ def on_input_disconnected(self, to_port, from_port):
         if not self.needCook:
             self.needCook = True
             return
-        self.cook()
+        self.update_stream()
 
     def checkPortType(self, to_port, from_port):
         # None type port can connect with any other type port
@@ -193,7 +181,7 @@ def set_property(self, name, value):
         super(AutoNode, self).set_property(name, value)
         self.set_port_type(name, type(value).__name__)
         if name in self.model.custom_properties.keys():
-            self.cook()
+            self.update_stream()
 
     def set_port_type(self, port, data_type: str):
         current_port = None
@@ -253,8 +241,9 @@ def set_disabled(self, mode=False):
         if mode:
             self.when_disabled()
             if self.graph is None or self.graph.auto_update:
-                self.update_streams()
-        self.cook()
+                self.update_stream()
+        else:
+            self.update_stream()
 
     def _close_error(self):
         self._error = False

example_auto_nodes/node_base/subgraph_node.py

@@ -1,7 +1,7 @@
 from .auto_node import AutoNode
 from NodeGraphQt import SubGraph
 import json
-from NodeGraphQt.base.utils import topological_sort
+from NodeGraphQt import topological_sort_by_down
 
 
 class SubGraphNode(AutoNode, SubGraph):
@@ -121,12 +121,14 @@ def run(self):
         else:
             start_nodes = self.sub_graph_input_nodes
 
-        nodes = topological_sort(start_nodes=start_nodes)
+        nodes = topological_sort_by_down(start_nodes=start_nodes)
 
         for node in nodes:
-            node.cook(stream=True)
+            if node.disabled():
+                node.when_disabled()
+            else:
+                node.cook()
             if node.error():
-                self.error(node.view.toolTip())
                 break
 
     def delete(self):