Browse Source

UI: Add Component Menu to Node-Add menu

Joseph Brandenburg 2 months ago
parent
commit
27763b12e0
3 changed files with 101 additions and 1 deletions
  1. 3 1
      __init__.py
  2. 46 0
      menu_classes.py
  3. 52 0
      ops_ui.py

+ 3 - 1
__init__.py

@@ -342,7 +342,7 @@ 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.
 
-from .menu_classes import node_context_menu_draw
+from .menu_classes import node_context_menu_draw, node_add_menu_draw
 
 def register():
     from bpy.utils import register_class
@@ -356,6 +356,8 @@ def register():
     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)
+    bpy.types.NODE_MT_add.append(node_add_menu_draw)
+
 
 
     km, kmi = init_keymaps()

+ 46 - 0
menu_classes.py

@@ -3,6 +3,7 @@ from bpy.types import Panel, Menu
 def TellClasses():
     return [
         MantisActiveTreePanel,
+        MantisNodeGroupsMenu,
     ]
 
 
@@ -17,6 +18,51 @@ def node_context_menu_draw(self, context):
     layout.operator("mantis.import_from_component_library")
     # layout.menu('NODE_MT_context_menu_mantis')
 
+
+# Function to append submenu to node add menu
+def node_add_menu_draw(self, context):
+    # NODE_MT_add
+    layout = self.layout
+    layout.separator()  # Optional: Adds a separator before your submenu
+    layout.menu("NODE_MT_add_mantis_groups")
+
+    # layout.menu('NODE_MT_context_menu_mantis')
+
+class MantisNodeGroupsMenu(Menu):
+    """Menu to show available node groups"""
+    bl_idname= "NODE_MT_add_mantis_groups"
+    bl_label = "Components"
+    def draw(self, context):
+        node_tree = None
+        # just gonna do the same thing we do in poll operators
+        if not context.space_data:
+            return
+        if not hasattr(context.space_data, "path"):
+            return
+        try:
+            node_tree = context.space_data.path[0].node_tree
+        except IndexError: # not in the UI, for example, in a script instead.
+            return
+        if node_tree is None: # because the space is right but there is no selected tree.
+            return
+        # now we're clear to do the menu function
+        layout = self.layout
+        from bpy import data
+        for ng in data.node_groups:
+            if ng.bl_idname in ['MantisTree', 'SchemaTree']:
+                if ng.name == node_tree.name and ng.bl_idname == node_tree.bl_idname:
+                    continue # don't add the node group into itself.
+                if ng.contains_tree(node_tree):
+                    continue # don't create an infinite loop of node trees
+                operator_settings = layout.operator(
+                    "mantis.add_component", text=ng.name,
+                    icon='NODE', emboss=False,)
+                operator_settings.node_group_tree_name=ng.name
+                operator_settings.tree_invoked = node_tree.name
+
+
+
+
 class MantisActiveTreePanel(Panel):
     """N-Panel menu for Mantis"""
     bl_label = "Active Tree"

+ 52 - 0
ops_ui.py

@@ -10,10 +10,62 @@ from .utilities import (prRed, prGreen, prPurple, prWhite,
 
 def TellClasses():
     return [
+        MantisAddComponent,
         ColorPalleteAddSocket,
         ColorPalleteRemoveSocket,
     ]
 
+def poll_in_mantis_tree(context):
+    if not context.space_data:
+        return False
+    if not hasattr(context.space_data, "path"):
+        return False
+    try:
+        node_tree = context.space_data.path[0].node_tree
+    except IndexError: # not in the UI, for example, in a script instead.
+        return False
+    if node_tree is None: # because the space is right but there is no selected tree.
+        return False
+    return True
+
+class MantisAddComponent(Operator):
+    """Add a component group into the scene"""
+    bl_idname = "mantis.add_component"
+    bl_label = "Add Component Group"
+    bl_options = {'REGISTER', 'UNDO'}
+
+    node_group_tree_name : bpy.props.StringProperty()
+    tree_invoked : bpy.props.StringProperty(options ={'HIDDEN'})
+
+    @classmethod
+    def poll(cls, context):
+        return poll_in_mantis_tree(context)
+
+    def invoke(self, context, event):
+        return self.execute(context)
+    
+    def execute (self, context):
+        from bpy import data
+        tree = bpy.data.node_groups[self.tree_invoked]
+        for node in tree.nodes:
+            node.select = False
+        tree = bpy.data.node_groups[self.tree_invoked]
+        node_group = data.node_groups.get(self.node_group_tree_name)
+        node_group_type = "MantisNodeGroup"
+        if node_group.bl_idname == 'SchemaTree':
+            node_group_type = "MantisSchemaGroup"
+        node = tree.nodes.new(node_group_type)
+        node.select = True
+        tree.nodes.active = node
+        node.node_tree = node_group
+        node.location = context.space_data.cursor_location
+        return bpy.ops.node.translate_attach_remove_on_cancel('INVOKE_DEFAULT',)
+                # TRANSFORM_OT_translate={ "value":
+                #                         (context.space_data.cursor_location[0]*1.25,
+                #                          context.space_data.cursor_location[1]*1.25, 0)})
+# TODO: figure out how to make this feel nice in the UI
+# bpy.ops.node.translate_attach_remove_on_cancel(TRANSFORM_OT_translate={"value":(-59.247, -5.7344, 0), "orient_type":'GLOBAL', "orient_matrix":((1, 0, 0), (0, 1, 0), (0, 0, 1)), "orient_matrix_type":'GLOBAL', "constraint_axis":(False, False, False), "mirror":False, "use_proportional_edit":False, "proportional_edit_falloff":'SMOOTH', "proportional_size":1, "use_proportional_connected":False, "use_proportional_projected":False, "snap":False, "snap_elements":{'GRID'}, "use_snap_project":False, "snap_target":'CLOSEST', "use_snap_self":True, "use_snap_edit":True, "use_snap_nonedit":True, "use_snap_selectable":False, "snap_point":(0, 0, 0), "snap_align":False, "snap_normal":(0, 0, 0), "gpencil_strokes":False, "cursor_transform":False, "texture_space":False, "remove_on_cancel":True, "use_duplicated_keyframes":False, "view2d_edge_pan":True, "release_confirm":False, "use_accurate":False, "use_automerge_and_split":False, "translate_origin":False}, NODE_OT_attach={})
+
 class ColorPalleteAddSocket(Operator):
     """Add a Color Set socket to the node"""
     bl_idname = "mantis.color_pallete_socket_add"