node_wrangler/interface.py

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