add : node menu can add by node class, the command can alse inherit by class relationship

Author: ArnoChenFx

NodeGraphQt/base/menu.py

@@ -161,24 +161,32 @@ class NodesMenu(NodeGraphMenu):
         nodes_menu = node_graph.get_context_menu('nodes')
     """
 
-    def add_command(self, name, func=None, node_type=None):
+    def add_command(self, name, func=None, node_type=None, node_class=None):
         """
         Re-implemented to add a command to the specified node type menu.
 
         Args:
             name (str): command name.
             func (function): command function eg. "func(``graph``, ``node``)".
             node_type (str): specified node type for the command.
+            node_class (class): specified node class for the command.
 
         Returns:
             NodeGraphQt.NodeGraphCommand: the appended command.
         """
-        if not node_type:
-            raise NodeMenuError('Node type not specified!')
+        if not node_type and not node_class:
+            raise NodeMenuError('Node type or Node class not specified!')
+        if node_class:
+            node_type = node_class.__name__
 
         node_menu = self.qmenu.get_menu(node_type)
         if not node_menu:
             node_menu = BaseMenu(node_type, self.qmenu)
+
+            if node_class:
+                node_menu.node_class = node_class
+                node_menu.graph = self._graph
+
             self.qmenu.addMenu(node_menu)
 
         if not self.qmenu.isEnabled():
@@ -190,6 +198,14 @@ def add_command(self, name, func=None, node_type=None):
             action.setShortcutVisibleInContextMenu(True)
         if func:
             action.executed.connect(func)
+
+        if node_class:
+            node_menus = self.qmenu.get_menus(node_class)
+            if node_menu in node_menus:
+                node_menus.remove(node_menu)
+            for menu in node_menus:
+                menu.addAction(action)
+
         qaction = node_menu.addAction(action)
         return NodeGraphCommand(self._graph, qaction)
 

NodeGraphQt/widgets/actions.py

@@ -8,6 +8,8 @@ class BaseMenu(QtWidgets.QMenu):
     def __init__(self, *args, **kwargs):
         super(BaseMenu, self).__init__(*args, **kwargs)
         self.setStyleSheet(STYLE_QMENU)
+        self.node_class = None
+        self.graph = None
 
     # disable for issue #142
     # def hideEvent(self, event):
@@ -16,11 +18,26 @@ def __init__(self, *args, **kwargs):
     #         if hasattr(a, 'node_id'):
     #             a.node_id = None
 
-    def get_menu(self, name):
+    def get_menu(self, name,node_id=None):
         for action in self.actions():
-            if action.menu() and action.menu().title() == name:
-                return action.menu()
-
+            menu = action.menu()
+            if not menu:
+                continue
+            if menu.title() == name:
+                return menu
+            if node_id and menu.node_class:
+                node = menu.graph.get_node_by_id(node_id)
+                if isinstance(node,menu.node_class):
+                    return menu
+
+    def get_menus(self,node_class):
+        menus = []
+        for action in self.actions():
+            menu = action.menu()
+            if menu.node_class:
+                if issubclass(menu.node_class,node_class):
+                    menus.append(menu)
+        return menus
 
 class GraphAction(QtWidgets.QAction):
 

NodeGraphQt/widgets/viewer.py

@@ -167,7 +167,7 @@ def contextMenuEvent(self, event):
             nodes = [i for i in items if isinstance(i, AbstractNodeItem)]
             if nodes:
                 node = nodes[0]
-                ctx_menu = self._ctx_node_menu.get_menu(node.type_)
+                ctx_menu = self._ctx_node_menu.get_menu(node.type_,node.id)
                 if ctx_menu:
                     for action in ctx_menu.actions():
                         if not action.menu():

example_auto_nodes/input_nodes.py

@@ -1,5 +1,5 @@
 from NodeGraphQt import QtCore
-from.node_base.auto_node import AutoNode
+from .node_base.auto_node import AutoNode
 import os
 
 class FloatInputNode(AutoNode):

example_auto_nodes/module_nodes.py

@@ -1,6 +1,6 @@
-from example_auto_nodes.node_base.module_node import (ModuleNode,
-                                                      get_functions_from_module,
-                                                      get_functions_from_type)
+from .node_base.module_node import (ModuleNode,
+                                    get_functions_from_module,
+                                    get_functions_from_type)
 
 import math
 import os
@@ -31,7 +31,7 @@ class MathModuleNode(ModuleNode):
     module_functions = get_functions_from_type(math)
 
     def __init__(self):
-        super(MathModuleNode, self).__init__()
+        super(MathModuleNode, self).__init__(float, float)
         self.defaultValue = 0.0
 
 
@@ -52,6 +52,7 @@ class osModuleNode(ModuleNode):
     def __init__(self):
         super(osModuleNode, self).__init__()
 
+
 class sysModuleNode(ModuleNode):
     """
     sys module node.
@@ -82,8 +83,7 @@ class randomModuleNode(ModuleNode):
     # set the initial default node name.
     NODE_NAME = 'random module'
 
-    module_functions = get_functions_from_module(random,max_depth=2)
-
+    module_functions = get_functions_from_module(random, max_depth=2)
 
     def __init__(self):
         super(randomModuleNode, self).__init__()
@@ -100,21 +100,20 @@ class numpyModuleNode(ModuleNode):
     # set the initial default node name.
     NODE_NAME = 'numpy module'
 
-    module_functions = get_functions_from_module(numpy,max_depth=2)
-
+    module_functions = get_functions_from_module(numpy, max_depth=2)
 
     def __init__(self):
         super(numpyModuleNode, self).__init__()
 
-    def is_function(self,obj):
-        result = super(numpyModuleNode,self).is_function(obj)
+    def is_function(self, obj):
+        result = super(numpyModuleNode, self).is_function(obj)
         if result:
             return True
         elif type(obj).__name__ == "ufunc":
             return True
         return False
 
-    def get_numpy_args(self,func):
+    def get_numpy_args(self, func):
         args = []
         info = numpydoc.docscrape.FunctionDoc(func)
         for i in info["Parameters"]:
@@ -139,7 +138,7 @@ def addFunction(self, prop, func):
                     args = self.get_numpy_args(self.func)
                 except Exception as error:
                     if type(self.func).__name__ == "ufunc":
-                        args = ["input"+str(i+1) for i in range(self.func.nin)]
+                        args = ["input" + str(i + 1) for i in range(self.func.nin)]
 
         self.process_args(args)
 

example_auto_nodes/node_base/auto_node.py

@@ -16,14 +16,17 @@ def rand_color(seed_type):
 
 
 class AutoNode(BaseNode):
-    def __init__(self):
+    def __init__(self,defaultInputType=None,defaultOutputType=None):
         super(AutoNode, self).__init__()
         self.needCook = True
-        self.defaultValue = None
         self._error = False
+        self.matchTypes = [[float, int]]
+        self.errorColor = (200, 50, 50)
 
         self.defaultColor = self.get_property("color")
-        self.errorColor = (200,50,50)
+        self.defaultValue = None
+        self.defaultInputType = defaultInputType
+        self.defaultOutputType = defaultOutputType
 
     def cookNextNode(self):
         for nodeList in self.connected_output_nodes().values():
@@ -93,10 +96,12 @@ def set_disabled(self, mode=False):
         if self.input_ports():
             self.cook()
 
-    @staticmethod
-    def checkPortType(to_port, from_port):
+    def checkPortType(self, to_port, from_port):
         if hasattr(to_port,"DataType") and hasattr(from_port,"DataType"):
             if to_port.DataType is not from_port.DataType:
+                for types in self.matchTypes:
+                    if to_port.DataType in types and from_port.DataType in types:
+                        return True
                 return False
 
         return True
@@ -106,10 +111,6 @@ def set_property(self, name, value):
         self.set_port_type(name, type(value))
 
     def set_port_type(self, port, value_type):
-
-        if value_type is int:
-            value_type = float
-
         current_port = None
 
         if type(port) is Port:
@@ -147,13 +148,17 @@ def add_input(self, name='input', data_type=None, multi_input=False, display_nam
         new_port = super(AutoNode, self).add_input(name, multi_input, display_name, color)
         if data_type:
             self.set_port_type(new_port, data_type)
+        elif self.defaultInputType:
+            self.set_port_type(new_port, self.defaultInputType)
         return new_port
 
     def add_output(self, name='output', data_type=None, multi_output=True, display_name=True,
                    color=None):
         new_port = super(AutoNode, self).add_output(name, multi_output, display_name, color)
         if data_type:
             self.set_port_type(new_port, data_type)
+        elif self.defaultOutputType:
+            self.set_port_type(new_port, self.defaultOutputType)
         return new_port
 
     def _close_error(self):
@@ -162,6 +167,10 @@ def _close_error(self):
         self._view._tooltip_disable(False)
 
     def _show_error(self,message):
+        if not self._error:
+            self.defaultColor = self.get_property("color")
+
+        self._error = True
         self.set_property('color', self.errorColor)
         tooltip = '<b>{}</b>'.format(self.name())
         tooltip += ' <font color="red"><br>({})</br></font>'.format(message)
@@ -172,5 +181,4 @@ def error(self,message = None):
         if message is None:
             return self._error
 
-        self._error = True
         self._show_error(message)

example_auto_nodes/node_base/module_node.py

@@ -1,5 +1,5 @@
 import inspect
-from example_auto_nodes.node_base.auto_node import AutoNode
+from .auto_node import AutoNode
 
 
 def _get_functions_from_module(module, function_dict, max_depth=1, module_name=None):
@@ -44,8 +44,8 @@ class ModuleNode(AutoNode):
 
     module_functions = {}
 
-    def __init__(self):
-        super(ModuleNode, self).__init__()
+    def __init__(self,defaultInputType=None,defaultOutputType=None):
+        super(ModuleNode, self).__init__(defaultInputType,defaultOutputType)
         self.set_color(25, 58, 51)
         self.add_combo_menu('funcs', 'Functions', items=list(self.module_functions.keys()), tab='widgets')