Add Operators to Context Menu

__init__.py

@@ -331,6 +331,16 @@ def on_undo_post_handler(scene): # the undo will trigger a depsgraph update
             # set the tree to invalid to trigger a tree update
             # since the context data is wiped by an undo.
 
+# Function to append submenu to context menu
+def node_context_menu_draw(self, context):
+    layout = self.layout
+    layout.separator()  # Optional: Adds a separator before your submenu
+    layout.operator_context = "INVOKE_DEFAULT"
+    layout.operator("mantis.nodes_cleanup", text='Sort Selected Nodes')
+    layout.operator("mantis.connect_nodes_to_input")
+    layout.operator("mantis.select_nodes_of_type")
+    # layout.menu('NODE_MT_context_menu_mantis')
+
 def register():
     from bpy.utils import register_class
     
@@ -342,6 +352,7 @@ def register():
             raise e
     nodeitems_utils.register_node_categories('MantisNodeCategories', node_categories)
     nodeitems_utils.register_node_categories('SchemaNodeCategories', schema_categories)
+    bpy.types.NODE_MT_context_menu.append(node_context_menu_draw)
 
 
     km, kmi = init_keymaps()
@@ -361,6 +372,7 @@ def register():
 
 
 def unregister():
+    bpy.types.NODE_MT_context_menu.remove(node_context_menu_draw)
     for tree in bpy.data.node_groups: # ensure it doesn't try to update while quitting.
         if tree.bl_idname in ['MantisTree, SchemaTree']:
             tree.is_exporting=True; tree.is_executing=True

menu_classes.py

@@ -1,4 +1,4 @@
-from bpy.types import Panel
+from bpy.types import Panel, Menu
 
 def TellClasses():
     return [
@@ -35,3 +35,4 @@ class MantisActiveTreePanel(Panel):
         layout.prop(nt, "do_live_update", text='Live Updates',)
         layout.operator("mantis.invalidate_node_tree")
         layout.operator("mantis.execute_node_tree")
+        layout.operator("mantis.force_display_update", text='Force Display Update')

ops_nodegroup.py

@@ -349,21 +349,41 @@ class ConnectNodeToInput(Operator):
         name="Node Input Socket",
         description="Select which of this node's sockets to recieve the connection",)
     tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
+    world_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
+    material_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
     node_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
 
     @classmethod
     def poll(cls, context):
-        return (any_tree_poll(context))
+        is_tree = (any_tree_poll(context))
+        is_material = False
+        if hasattr(context, 'space_data') and context.space_data.type == 'NODE_EDITOR':
+            is_material = context.space_data.node_tree.bl_idname == 'ShaderNodeTree'
+        return is_tree and not is_material # doesn't work for Material right now TODO
 
     def invoke(self, context, event):
+        self.material_invoked=''; self.world_invoked=''
         self.tree_invoked = context.active_node.id_data.name
+        if context.space_data.node_tree.bl_idname == 'ShaderNodeTree':
+            if context.space_data.shader_type == 'WORLD':
+                self.world_invoked = context.space_data.id.name
+            elif context.space_data.shader_type == 'OBJECT':
+                self.material_invoked = context.space_data.id.name
+            else:
+                return {'CANCELLED'} # what is the LINESTYLE menu? TODO
         self.node_invoked = context.active_node.name
         # we use active_node here ^ because we are comparing the active node to the selection.
         wm = context.window_manager
         return wm.invoke_props_dialog(self)
     
     def execute(self, context):
-        t = bpy.data.node_groups[self.tree_invoked]
+        if context.space_data.node_tree.bl_idname == 'ShaderNodeTree':
+            if context.space_data.shader_type == 'WORLD':
+                t = bpy.data.worlds[self.material_invoked].node_tree
+            elif context.space_data.shader_type == 'OBJECT':
+                t = bpy.data.materials[self.material_invoked].node_tree
+        else:
+            t = bpy.data.node_groups.get(self.tree_invoked)
         if hasattr(t, "is_executing"): # for Mantis trees, but this function should just work anywhere.
             t.is_executing = True
         n = t.nodes[self.node_invoked]
@@ -384,6 +404,9 @@ class ConnectNodeToInput(Operator):
                         else: connect_me = s
                         inp.location = node.location
                         inp.location.x-=200
+                    if connect_me is None:
+                        continue # I don't know how this happens.
+                    # TODO: this bit doesn't work in shader trees for some reason, fix it
                     t.links.new(input=connect_me, output=connect_to_me)
 
         if hasattr(t, "is_executing"):