Geometry Nodes: improve operators for node editor header

This allows users to create new modifiers directly from the
Geometry Nodes Editor.

2 files changed, 62 insertions, 15 deletions

diff --git a/release/scripts/startup/bl_operators/geometry_nodes.py b/release/scripts/startup/bl_operators/geometry_nodes.py
index 374bd4161a7..23926f28ab0 100644
--- a/release/scripts/startup/bl_operators/geometry_nodes.py
+++ b/release/scripts/startup/bl_operators/geometry_nodes.py
@@ -18,6 +18,20 @@
 
 import bpy
 
+def geometry_node_group_empty_new(context):
+    group = bpy.data.node_groups.new("Geometry Nodes", 'GeometryNodeTree')
+    group.inputs.new('NodeSocketGeometry', "Geometry")
+    group.outputs.new('NodeSocketGeometry', "Geometry")
+    input_node = group.nodes.new('NodeGroupInput')
+    output_node = group.nodes.new('NodeGroupOutput')
+    output_node.is_active_output = True
+
+    input_node.location.x = -200 - input_node.width
+    output_node.location.x = 200
+
+    group.links.new(output_node.inputs[0], input_node.outputs[0])
+
+    return group
 
 def geometry_modifier_poll(context) -> bool:
     ob = context.object
@@ -28,11 +42,34 @@ def geometry_modifier_poll(context) -> bool:
 
     return True
 
+class NewGeometryNodeModifier(bpy.types.Operator):
+    """Create a new modifier with a new geometry node group"""
+
+    bl_idname = "node.new_geometry_node_modifier"
+    bl_label = "New Geometry Node Modifier"
+    bl_options = {'REGISTER', 'UNDO'}
+
+    @classmethod
+    def poll(cls, context):
+        return geometry_modifier_poll(context)
+
+    def execute(self, context):
+        modifier = context.object.modifiers.new("Empty", "NODES")
+
+        if not modifier:
+            return {'CANCELLED'}
+
+        group = geometry_node_group_empty_new(context)
+        modifier.node_group = group
+
+        return {'FINISHED'}
+
+
+class NewGeometryNodeTreeAssign(bpy.types.Operator):
+    """Create a new geometry node group and assign it the the active modifier"""
 
-class NewGeometryNodeTree(bpy.types.Operator):
-    """Create a new geometry node tree"""
-    bl_idname = "node.new_geometry_node_tree"
-    bl_label = "New Geometry Node Tree"
+    bl_idname = "node.new_geometry_node_group_assign"
+    bl_label = "Assign New Geometry Node Group"
     bl_options = {'REGISTER', 'UNDO'}
 
     @classmethod
@@ -40,21 +77,18 @@ class NewGeometryNodeTree(bpy.types.Operator):
         return geometry_modifier_poll(context)
 
     def execute(self, context):
-        group = bpy.data.node_groups.new("Geometry Nodes", 'GeometryNodeTree')
-        group.inputs.new('NodeSocketGeometry', "Geometry")
-        group.outputs.new('NodeSocketGeometry', "Geometry")
-        input_node = group.nodes.new('NodeGroupInput')
-        output_node = group.nodes.new('NodeGroupOutput')
-        output_node.is_active_output = True
+        modifier = context.object.modifiers.active
 
-        input_node.location.x = -200 - input_node.width
-        output_node.location.x = 200
+        if not modifier:
+            return {'CANCELLED'}
 
-        group.links.new(output_node.inputs[0], input_node.outputs[0])
+        group = geometry_node_group_empty_new(context)
+        modifier.node_group = group
 
         return {'FINISHED'}
 
 
 classes = (
-    NewGeometryNodeTree,
+    NewGeometryNodeModifier,
+    NewGeometryNodeTreeAssign,
 )
diff --git a/release/scripts/startup/bl_ui/space_node.py b/release/scripts/startup/bl_ui/space_node.py
index 4cd38831cdf..9ad454d49e1 100644
--- a/release/scripts/startup/bl_ui/space_node.py
+++ b/release/scripts/startup/bl_ui/space_node.py
@@ -154,7 +154,20 @@ class NODE_HT_header(Header):
         elif snode.tree_type == 'GeometryNodeTree':
             NODE_MT_editor_menus.draw_collapsible(context, layout)
             layout.separator_spacer()
-            layout.template_ID(snode, "node_tree", new="node.new_geometry_node_tree")
+
+            ob = context.object
+
+            row = layout.row()
+            if snode.pin:
+                row.enabled = False
+                row.template_ID(snode, "node_tree", new="node.new_geometry_node_group_assign")
+            elif ob:
+                active_modifier = ob.modifiers.active
+                if active_modifier and active_modifier.type == "NODES":                     
+                    row.template_ID(active_modifier, "node_group", new="node.new_geometry_node_group_assign")
+                else:
+                    row.template_ID(snode, "node_tree", new="node.new_geometry_node_modifier")
+            
 
         else:
             # Custom node tree is edited as independent ID block