node_wrangler/interface.py

   1 # SPDX-FileCopyrightText: 2023 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import bpy
   6 from bpy.types import Panel, Menu
   7 from bpy.props import StringProperty
   8 from nodeitems_utils import node_categories_iter, NodeItemCustom
   9
  10 from . import operators
  11
  12 from .utils.constants import blend_types, geo_combine_operations, operations
  13 from .utils.nodes import get_nodes_links, nw_check, NWBase
  14
  15
  16 def drawlayout(context, layout, mode='non-panel'):
  17     tree_type = context.space_data.tree_type
  18
  19     col = layout.column(align=True)
  20     col.menu(NWMergeNodesMenu.bl_idname)
  21     col.separator()
  22
  23     col = layout.column(align=True)
  24     col.menu(NWSwitchNodeTypeMenu.bl_idname, text="Switch Node Type")
  25     col.separator()
  26
  27     if tree_type == 'ShaderNodeTree':
  28         col = layout.column(align=True)
  29         col.operator(operators.NWAddTextureSetup.bl_idname, text="Add Texture Setup", icon='NODE_SEL')
  30         col.operator(operators.NWAddPrincipledSetup.bl_idname, text="Add Principled Setup", icon='NODE_SEL')
  31         col.separator()
  32
  33     col = layout.column(align=True)
  34     col.operator(operators.NWDetachOutputs.bl_idname, icon='UNLINKED')
  35     col.operator(operators.NWSwapLinks.bl_idname)
  36     col.menu(NWAddReroutesMenu.bl_idname, text="Add Reroutes", icon='LAYER_USED')
  37     col.separator()
  38
  39     col = layout.column(align=True)
  40     col.menu(NWLinkActiveToSelectedMenu.bl_idname, text="Link Active To Selected", icon='LINKED')
  41     if tree_type != 'GeometryNodeTree':
  42         col.operator(operators.NWLinkToOutputNode.bl_idname, icon='DRIVER')
  43     col.separator()
  44
  45     col = layout.column(align=True)
  46     if mode == 'panel':
  47         row = col.row(align=True)
  48         row.operator(operators.NWClearLabel.bl_idname).option = True
  49         row.operator(operators.NWModifyLabels.bl_idname)
  50     else:
  51         col.operator(operators.NWClearLabel.bl_idname).option = True
  52         col.operator(operators.NWModifyLabels.bl_idname)
  53     col.menu(NWBatchChangeNodesMenu.bl_idname, text="Batch Change")
  54     col.separator()
  55     col.menu(NWCopyToSelectedMenu.bl_idname, text="Copy to Selected")
  56     col.separator()
  57
  58     col = layout.column(align=True)
  59     if tree_type == 'CompositorNodeTree':
  60         col.operator(operators.NWResetBG.bl_idname, icon='ZOOM_PREVIOUS')
  61     if tree_type != 'GeometryNodeTree':
  62         col.operator(operators.NWReloadImages.bl_idname, icon='FILE_REFRESH')
  63     col.separator()
  64
  65     col = layout.column(align=True)
  66     col.operator(operators.NWFrameSelected.bl_idname, icon='STICKY_UVS_LOC')
  67     col.separator()
  68
  69     col = layout.column(align=True)
  70     col.operator(operators.NWAlignNodes.bl_idname, icon='CENTER_ONLY')
  71     col.separator()
  72
  73     col = layout.column(align=True)
  74     col.operator(operators.NWDeleteUnused.bl_idname, icon='CANCEL')
  75     col.separator()
  76
  77
  78 class NodeWranglerPanel(Panel, NWBase):
  79     bl_idname = "NODE_PT_nw_node_wrangler"
  80     bl_space_type = 'NODE_EDITOR'
  81     bl_label = "Node Wrangler"
  82     bl_region_type = "UI"
  83     bl_category = "Node Wrangler"
  84
  85     prepend: StringProperty(
  86         name='prepend',
  87     )
  88     append: StringProperty()
  89     remove: StringProperty()
  90
  91     def draw(self, context):
  92         self.layout.label(text="(Quick access: Shift+W)")
  93         drawlayout(context, self.layout, mode='panel')
  94
  95
  96 #
  97 #  M E N U S
  98 #
  99 class NodeWranglerMenu(Menu, NWBase):
 100     bl_idname = "NODE_MT_nw_node_wrangler_menu"
 101     bl_label = "Node Wrangler"
 102
 103     def draw(self, context):
 104         self.layout.operator_context = 'INVOKE_DEFAULT'
 105         drawlayout(context, self.layout)
 106
 107
 108 class NWMergeNodesMenu(Menu, NWBase):
 109     bl_idname = "NODE_MT_nw_merge_nodes_menu"
 110     bl_label = "Merge Selected Nodes"
 111
 112     def draw(self, context):
 113         type = context.space_data.tree_type
 114         layout = self.layout
 115         if type == 'ShaderNodeTree':
 116             layout.menu(NWMergeShadersMenu.bl_idname, text="Use Shaders")
 117         if type == 'GeometryNodeTree':
 118             layout.menu(NWMergeGeometryMenu.bl_idname, text="Use Geometry Nodes")
 119             layout.menu(NWMergeMathMenu.bl_idname, text="Use Math Nodes")
 120         else:
 121             layout.menu(NWMergeMixMenu.bl_idname, text="Use Mix Nodes")
 122             layout.menu(NWMergeMathMenu.bl_idname, text="Use Math Nodes")
 123             props = layout.operator(operators.NWMergeNodes.bl_idname, text="Use Z-Combine Nodes")
 124             props.mode = 'MIX'
 125             props.merge_type = 'ZCOMBINE'
 126             props = layout.operator(operators.NWMergeNodes.bl_idname, text="Use Alpha Over Nodes")
 127             props.mode = 'MIX'
 128             props.merge_type = 'ALPHAOVER'
 129
 130
 131 class NWMergeGeometryMenu(Menu, NWBase):
 132     bl_idname = "NODE_MT_nw_merge_geometry_menu"
 133     bl_label = "Merge Selected Nodes using Geometry Nodes"
 134
 135     def draw(self, context):
 136         layout = self.layout
 137         # The boolean node + Join Geometry node
 138         for type, name, description in geo_combine_operations:
 139             props = layout.operator(operators.NWMergeNodes.bl_idname, text=name)
 140             props.mode = type
 141             props.merge_type = 'GEOMETRY'
 142
 143
 144 class NWMergeShadersMenu(Menu, NWBase):
 145     bl_idname = "NODE_MT_nw_merge_shaders_menu"
 146     bl_label = "Merge Selected Nodes using Shaders"
 147
 148     def draw(self, context):
 149         layout = self.layout
 150         for type in ('MIX', 'ADD'):
 151             name = f'{type.capitalize()} Shader'
 152             props = layout.operator(operators.NWMergeNodes.bl_idname, text=name)
 153             props.mode = type
 154             props.merge_type = 'SHADER'
 155
 156
 157 class NWMergeMixMenu(Menu, NWBase):
 158     bl_idname = "NODE_MT_nw_merge_mix_menu"
 159     bl_label = "Merge Selected Nodes using Mix"
 160
 161     def draw(self, context):
 162         layout = self.layout
 163         for type, name, description in blend_types:
 164             props = layout.operator(operators.NWMergeNodes.bl_idname, text=name)
 165             props.mode = type
 166             props.merge_type = 'MIX'
 167
 168
 169 class NWConnectionListOutputs(Menu, NWBase):
 170     bl_idname = "NODE_MT_nw_connection_list_out"
 171     bl_label = "From:"
 172
 173     def draw(self, context):
 174         layout = self.layout
 175         nodes, links = get_nodes_links(context)
 176
 177         n1 = nodes[context.scene.NWLazySource]
 178         for index, output in enumerate(n1.outputs):
 179             # Only show sockets that are exposed.
 180             if output.enabled:
 181                 layout.operator(
 182                     operators.NWCallInputsMenu.bl_idname,
 183                     text=output.name,
 184                     icon="RADIOBUT_OFF").from_socket = index
 185
 186
 187 class NWConnectionListInputs(Menu, NWBase):
 188     bl_idname = "NODE_MT_nw_connection_list_in"
 189     bl_label = "To:"
 190
 191     def draw(self, context):
 192         layout = self.layout
 193         nodes, links = get_nodes_links(context)
 194
 195         n2 = nodes[context.scene.NWLazyTarget]
 196
 197         for index, input in enumerate(n2.inputs):
 198             # Only show sockets that are exposed.
 199             # This prevents, for example, the scale value socket
 200             # of the vector math node being added to the list when
 201             # the mode is not 'SCALE'.
 202             if input.enabled:
 203                 op = layout.operator(operators.NWMakeLink.bl_idname, text=input.name, icon="FORWARD")
 204                 op.from_socket = context.scene.NWSourceSocket
 205                 op.to_socket = index
 206
 207
 208 class NWMergeMathMenu(Menu, NWBase):
 209     bl_idname = "NODE_MT_nw_merge_math_menu"
 210     bl_label = "Merge Selected Nodes using Math"
 211
 212     def draw(self, context):
 213         layout = self.layout
 214         for type, name, description in operations:
 215             props = layout.operator(operators.NWMergeNodes.bl_idname, text=name)
 216             props.mode = type
 217             props.merge_type = 'MATH'
 218
 219
 220 class NWBatchChangeNodesMenu(Menu, NWBase):
 221     bl_idname = "NODE_MT_nw_batch_change_nodes_menu"
 222     bl_label = "Batch Change Selected Nodes"
 223
 224     def draw(self, context):
 225         layout = self.layout
 226         layout.menu(NWBatchChangeBlendTypeMenu.bl_idname)
 227         layout.menu(NWBatchChangeOperationMenu.bl_idname)
 228
 229
 230 class NWBatchChangeBlendTypeMenu(Menu, NWBase):
 231     bl_idname = "NODE_MT_nw_batch_change_blend_type_menu"
 232     bl_label = "Batch Change Blend Type"
 233
 234     def draw(self, context):
 235         layout = self.layout
 236         for type, name, description in blend_types:
 237             props = layout.operator(operators.NWBatchChangeNodes.bl_idname, text=name)
 238             props.blend_type = type
 239             props.operation = 'CURRENT'
 240
 241
 242 class NWBatchChangeOperationMenu(Menu, NWBase):
 243     bl_idname = "NODE_MT_nw_batch_change_operation_menu"
 244     bl_label = "Batch Change Math Operation"
 245
 246     def draw(self, context):
 247         layout = self.layout
 248         for type, name, description in operations:
 249             props = layout.operator(operators.NWBatchChangeNodes.bl_idname, text=name)
 250             props.blend_type = 'CURRENT'
 251             props.operation = type
 252
 253
 254 class NWCopyToSelectedMenu(Menu, NWBase):
 255     bl_idname = "NODE_MT_nw_copy_node_properties_menu"
 256     bl_label = "Copy to Selected"
 257
 258     def draw(self, context):
 259         layout = self.layout
 260         layout.operator(operators.NWCopySettings.bl_idname, text="Settings from Active")
 261         layout.menu(NWCopyLabelMenu.bl_idname)
 262
 263
 264 class NWCopyLabelMenu(Menu, NWBase):
 265     bl_idname = "NODE_MT_nw_copy_label_menu"
 266     bl_label = "Copy Label"
 267
 268     def draw(self, context):
 269         layout = self.layout
 270         layout.operator(operators.NWCopyLabel.bl_idname, text="from Active Node's Label").option = 'FROM_ACTIVE'
 271         layout.operator(operators.NWCopyLabel.bl_idname, text="from Linked Node's Label").option = 'FROM_NODE'
 272         layout.operator(operators.NWCopyLabel.bl_idname, text="from Linked Output's Name").option = 'FROM_SOCKET'
 273
 274
 275 class NWAddReroutesMenu(Menu, NWBase):
 276     bl_idname = "NODE_MT_nw_add_reroutes_menu"
 277     bl_label = "Add Reroutes"
 278     bl_description = "Add Reroute Nodes to Selected Nodes' Outputs"
 279
 280     def draw(self, context):
 281         layout = self.layout
 282         layout.operator(operators.NWAddReroutes.bl_idname, text="to All Outputs").option = 'ALL'
 283         layout.operator(operators.NWAddReroutes.bl_idname, text="to Loose Outputs").option = 'LOOSE'
 284         layout.operator(operators.NWAddReroutes.bl_idname, text="to Linked Outputs").option = 'LINKED'
 285
 286
 287 class NWLinkActiveToSelectedMenu(Menu, NWBase):
 288     bl_idname = "NODE_MT_nw_link_active_to_selected_menu"
 289     bl_label = "Link Active to Selected"
 290
 291     def draw(self, context):
 292         layout = self.layout
 293         layout.menu(NWLinkStandardMenu.bl_idname)
 294         layout.menu(NWLinkUseNodeNameMenu.bl_idname)
 295         layout.menu(NWLinkUseOutputsNamesMenu.bl_idname)
 296
 297
 298 class NWLinkStandardMenu(Menu, NWBase):
 299     bl_idname = "NODE_MT_nw_link_standard_menu"
 300     bl_label = "To All Selected"
 301
 302     def draw(self, context):
 303         layout = self.layout
 304         props = layout.operator(operators.NWLinkActiveToSelected.bl_idname, text="Don't Replace Links")
 305         props.replace = False
 306         props.use_node_name = False
 307         props.use_outputs_names = False
 308         props = layout.operator(operators.NWLinkActiveToSelected.bl_idname, text="Replace Links")
 309         props.replace = True
 310         props.use_node_name = False
 311         props.use_outputs_names = False
 312
 313
 314 class NWLinkUseNodeNameMenu(Menu, NWBase):
 315     bl_idname = "NODE_MT_nw_link_use_node_name_menu"
 316     bl_label = "Use Node Name/Label"
 317
 318     def draw(self, context):
 319         layout = self.layout
 320         props = layout.operator(operators.NWLinkActiveToSelected.bl_idname, text="Don't Replace Links")
 321         props.replace = False
 322         props.use_node_name = True
 323         props.use_outputs_names = False
 324         props = layout.operator(operators.NWLinkActiveToSelected.bl_idname, text="Replace Links")
 325         props.replace = True
 326         props.use_node_name = True
 327         props.use_outputs_names = False
 328
 329
 330 class NWLinkUseOutputsNamesMenu(Menu, NWBase):
 331     bl_idname = "NODE_MT_nw_link_use_outputs_names_menu"
 332     bl_label = "Use Outputs Names"
 333
 334     def draw(self, context):
 335         layout = self.layout
 336         props = layout.operator(operators.NWLinkActiveToSelected.bl_idname, text="Don't Replace Links")
 337         props.replace = False
 338         props.use_node_name = False
 339         props.use_outputs_names = True
 340         props = layout.operator(operators.NWLinkActiveToSelected.bl_idname, text="Replace Links")
 341         props.replace = True
 342         props.use_node_name = False
 343         props.use_outputs_names = True
 344
 345
 346 class NWAttributeMenu(bpy.types.Menu):
 347     bl_idname = "NODE_MT_nw_node_attribute_menu"
 348     bl_label = "Attributes"
 349
 350     @classmethod
 351     def poll(cls, context):
 352         valid = False
 353         if nw_check(context):
 354             snode = context.space_data
 355             valid = snode.tree_type == 'ShaderNodeTree'
 356         return valid
 357
 358     def draw(self, context):
 359         l = self.layout
 360         nodes, links = get_nodes_links(context)
 361         mat = context.object.active_material
 362
 363         objs = []
 364         for obj in bpy.data.objects:
 365             for slot in obj.material_slots:
 366                 if slot.material == mat:
 367                     objs.append(obj)
 368         attrs = []
 369         for obj in objs:
 370             if obj.data.attributes:
 371                 for attr in obj.data.attributes:
 372                     attrs.append(attr.name)
 373         attrs = list(set(attrs))  # get a unique list
 374
 375         if attrs:
 376             for attr in attrs:
 377                 l.operator(operators.NWAddAttrNode.bl_idname, text=attr).attr_name = attr
 378         else:
 379             l.label(text="No attributes on objects with this material")
 380
 381
 382 class NWSwitchNodeTypeMenu(Menu, NWBase):
 383     bl_idname = "NODE_MT_nw_switch_node_type_menu"
 384     bl_label = "Switch Type to..."
 385
 386     def draw(self, context):
 387         layout = self.layout
 388         categories = [c for c in node_categories_iter(context)
 389                       if c.name not in ['Group', 'Script']]
 390         for cat in categories:
 391             idname = f"NODE_MT_nw_switch_{cat.identifier}_submenu"
 392             if hasattr(bpy.types, idname):
 393                 layout.menu(idname)
 394             else:
 395                 layout.label(text="Unable to load altered node lists.")
 396                 layout.label(text="Please re-enable Node Wrangler.")
 397                 break
 398
 399
 400 def draw_switch_category_submenu(self, context):
 401     layout = self.layout
 402     if self.category.name == 'Layout':
 403         for node in self.category.items(context):
 404             if node.nodetype != 'NodeFrame':
 405                 props = layout.operator(operators.NWSwitchNodeType.bl_idname, text=node.label)
 406                 props.to_type = node.nodetype
 407     else:
 408         for node in self.category.items(context):
 409             if isinstance(node, NodeItemCustom):
 410                 node.draw(self, layout, context)
 411                 continue
 412             props = layout.operator(operators.NWSwitchNodeType.bl_idname, text=node.label)
 413             props.to_type = node.nodetype
 414
 415 #
 416 #  APPENDAGES TO EXISTING UI
 417 #
 418
 419
 420 def select_parent_children_buttons(self, context):
 421     layout = self.layout
 422     layout.operator(operators.NWSelectParentChildren.bl_idname,
 423                     text="Select frame's members (children)").option = 'CHILD'
 424     layout.operator(operators.NWSelectParentChildren.bl_idname, text="Select parent frame").option = 'PARENT'
 425
 426
 427 def attr_nodes_menu_func(self, context):
 428     col = self.layout.column(align=True)
 429     col.menu("NODE_MT_nw_node_attribute_menu")
 430     col.separator()
 431
 432
 433 def multipleimages_menu_func(self, context):
 434     col = self.layout.column(align=True)
 435     col.operator(operators.NWAddMultipleImages.bl_idname, text="Multiple Images")
 436     col.operator(operators.NWAddSequence.bl_idname, text="Image Sequence")
 437     col.separator()
 438
 439
 440 def bgreset_menu_func(self, context):
 441     self.layout.operator(operators.NWResetBG.bl_idname)
 442
 443
 444 def save_viewer_menu_func(self, context):
 445     if nw_check(context):
 446         if context.space_data.tree_type == 'CompositorNodeTree':
 447             if context.scene.node_tree.nodes.active:
 448                 if context.scene.node_tree.nodes.active.type == "VIEWER":
 449                     self.layout.operator(operators.NWSaveViewer.bl_idname, icon='FILE_IMAGE')
 450
 451
 452 def reset_nodes_button(self, context):
 453     node_active = context.active_node
 454     node_selected = context.selected_nodes
 455     node_ignore = ["FRAME", "REROUTE", "GROUP"]
 456
 457     # Check if active node is in the selection and respective type
 458     if (len(node_selected) == 1) and node_active and node_active.select and node_active.type not in node_ignore:
 459         row = self.layout.row()
 460         row.operator(operators.NWResetNodes.bl_idname, text="Reset Node", icon="FILE_REFRESH")
 461         self.layout.separator()
 462
 463     elif (len(node_selected) == 1) and node_active and node_active.select and node_active.type == "FRAME":
 464         row = self.layout.row()
 465         row.operator(operators.NWResetNodes.bl_idname, text="Reset Nodes in Frame", icon="FILE_REFRESH")
 466         self.layout.separator()
 467
 468
 469 classes = (
 470     NodeWranglerPanel,
 471     NodeWranglerMenu,
 472     NWMergeNodesMenu,
 473     NWMergeGeometryMenu,
 474     NWMergeShadersMenu,
 475     NWMergeMixMenu,
 476     NWConnectionListOutputs,
 477     NWConnectionListInputs,
 478     NWMergeMathMenu,
 479     NWBatchChangeNodesMenu,
 480     NWBatchChangeBlendTypeMenu,
 481     NWBatchChangeOperationMenu,
 482     NWCopyToSelectedMenu,
 483     NWCopyLabelMenu,
 484     NWAddReroutesMenu,
 485     NWLinkActiveToSelectedMenu,
 486     NWLinkStandardMenu,
 487     NWLinkUseNodeNameMenu,
 488     NWLinkUseOutputsNamesMenu,
 489     NWAttributeMenu,
 490     NWSwitchNodeTypeMenu,
 491 )
 492
 493
 494 def register():
 495     from bpy.utils import register_class
 496     for cls in classes:
 497         register_class(cls)
 498
 499     # menu items
 500     bpy.types.NODE_MT_select.append(select_parent_children_buttons)
 501     bpy.types.NODE_MT_category_shader_input.prepend(attr_nodes_menu_func)
 502     bpy.types.NODE_PT_backdrop.append(bgreset_menu_func)
 503     bpy.types.NODE_PT_active_node_generic.append(save_viewer_menu_func)
 504     bpy.types.NODE_MT_category_shader_texture.prepend(multipleimages_menu_func)
 505     bpy.types.NODE_MT_category_compositor_input.prepend(multipleimages_menu_func)
 506     bpy.types.NODE_PT_active_node_generic.prepend(reset_nodes_button)
 507     bpy.types.NODE_MT_node.prepend(reset_nodes_button)
 508
 509
 510 def unregister():
 511     # menu items
 512     bpy.types.NODE_MT_select.remove(select_parent_children_buttons)
 513     bpy.types.NODE_MT_category_shader_input.remove(attr_nodes_menu_func)
 514     bpy.types.NODE_PT_backdrop.remove(bgreset_menu_func)
 515     bpy.types.NODE_PT_active_node_generic.remove(save_viewer_menu_func)
 516     bpy.types.NODE_MT_category_shader_texture.remove(multipleimages_menu_func)
 517     bpy.types.NODE_MT_category_compositor_input.remove(multipleimages_menu_func)
 518     bpy.types.NODE_PT_active_node_generic.remove(reset_nodes_button)
 519     bpy.types.NODE_MT_node.remove(reset_nodes_button)
 520
 521     from bpy.utils import unregister_class
 522     for cls in classes:
 523         unregister_class(cls)