node_wrangler/utils/nodes.py

   1 # SPDX-FileCopyrightText: 2023 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import bpy
   6 from bpy_extras.node_utils import connect_sockets
   7 from math import hypot, inf
   8
   9
  10 def force_update(context):
  11     context.space_data.node_tree.update_tag()
  12
  13
  14 def dpi_fac():
  15     prefs = bpy.context.preferences.system
  16     return prefs.dpi / 72
  17
  18
  19 def prefs_line_width():
  20     prefs = bpy.context.preferences.system
  21     return prefs.pixel_size
  22
  23
  24 def node_mid_pt(node, axis):
  25     if axis == 'x':
  26         d = node.location.x + (node.dimensions.x / 2)
  27     elif axis == 'y':
  28         d = node.location.y - (node.dimensions.y / 2)
  29     else:
  30         d = 0
  31     return d
  32
  33
  34 def autolink(node1, node2, links):
  35     available_inputs = [inp for inp in node2.inputs if inp.enabled]
  36     available_outputs = [outp for outp in node1.outputs if outp.enabled]
  37     for outp in available_outputs:
  38         for inp in available_inputs:
  39             if not inp.is_linked and inp.name == outp.name:
  40                 connect_sockets(outp, inp)
  41                 return True
  42
  43     for outp in available_outputs:
  44         for inp in available_inputs:
  45             if not inp.is_linked and inp.type == outp.type:
  46                 connect_sockets(outp, inp)
  47                 return True
  48
  49     # force some connection even if the type doesn't match
  50     if available_outputs:
  51         for inp in available_inputs:
  52             if not inp.is_linked:
  53                 connect_sockets(available_outputs[0], inp)
  54                 return True
  55
  56     # even if no sockets are open, force one of matching type
  57     for outp in available_outputs:
  58         for inp in available_inputs:
  59             if inp.type == outp.type:
  60                 connect_sockets(outp, inp)
  61                 return True
  62
  63     # do something!
  64     for outp in available_outputs:
  65         for inp in available_inputs:
  66             connect_sockets(outp, inp)
  67             return True
  68
  69     print("Could not make a link from " + node1.name + " to " + node2.name)
  70     return False
  71
  72
  73 def abs_node_location(node):
  74     abs_location = node.location
  75     if node.parent is None:
  76         return abs_location
  77     return abs_location + abs_node_location(node.parent)
  78
  79
  80 def node_at_pos(nodes, context, event):
  81     nodes_under_mouse = []
  82     target_node = None
  83
  84     store_mouse_cursor(context, event)
  85     x, y = context.space_data.cursor_location
  86
  87     # Make a list of each corner (and middle of border) for each node.
  88     # Will be sorted to find nearest point and thus nearest node
  89     node_points_with_dist = []
  90     for node in nodes:
  91         skipnode = False
  92         if node.type != 'FRAME':  # no point trying to link to a frame node
  93             dimx = node.dimensions.x / dpi_fac()
  94             dimy = node.dimensions.y / dpi_fac()
  95             locx, locy = abs_node_location(node)
  96
  97             if not skipnode:
  98                 node_points_with_dist.append([node, hypot(x - locx, y - locy)])  # Top Left
  99                 node_points_with_dist.append([node, hypot(x - (locx + dimx), y - locy)])  # Top Right
 100                 node_points_with_dist.append([node, hypot(x - locx, y - (locy - dimy))])  # Bottom Left
 101                 node_points_with_dist.append([node, hypot(x - (locx + dimx), y - (locy - dimy))])  # Bottom Right
 102
 103                 node_points_with_dist.append([node, hypot(x - (locx + (dimx / 2)), y - locy)])  # Mid Top
 104                 node_points_with_dist.append([node, hypot(x - (locx + (dimx / 2)), y - (locy - dimy))])  # Mid Bottom
 105                 node_points_with_dist.append([node, hypot(x - locx, y - (locy - (dimy / 2)))])  # Mid Left
 106                 node_points_with_dist.append([node, hypot(x - (locx + dimx), y - (locy - (dimy / 2)))])  # Mid Right
 107
 108     nearest_node = sorted(node_points_with_dist, key=lambda k: k[1])[0][0]
 109
 110     for node in nodes:
 111         if node.type != 'FRAME' and skipnode == False:
 112             locx, locy = abs_node_location(node)
 113             dimx = node.dimensions.x / dpi_fac()
 114             dimy = node.dimensions.y / dpi_fac()
 115             if (locx <= x <= locx + dimx) and \
 116                (locy - dimy <= y <= locy):
 117                 nodes_under_mouse.append(node)
 118
 119     if len(nodes_under_mouse) == 1:
 120         if nodes_under_mouse[0] != nearest_node:
 121             target_node = nodes_under_mouse[0]  # use the node under the mouse if there is one and only one
 122         else:
 123             target_node = nearest_node  # else use the nearest node
 124     else:
 125         target_node = nearest_node
 126     return target_node
 127
 128
 129 def store_mouse_cursor(context, event):
 130     space = context.space_data
 131     v2d = context.region.view2d
 132     tree = space.edit_tree
 133
 134     # convert mouse position to the View2D for later node placement
 135     if context.region.type == 'WINDOW':
 136         space.cursor_location_from_region(event.mouse_region_x, event.mouse_region_y)
 137     else:
 138         space.cursor_location = tree.view_center
 139
 140
 141 def get_nodes_links(context):
 142     tree = context.space_data.edit_tree
 143     return tree.nodes, tree.links
 144
 145
 146 viewer_socket_name = "tmp_viewer"
 147
 148
 149 def is_viewer_socket(socket):
 150     # checks if a internal socket is a valid viewer socket
 151     return socket.name == viewer_socket_name and socket.NWViewerSocket
 152
 153
 154 def get_internal_socket(socket):
 155     # get the internal socket from a socket inside or outside the group
 156     node = socket.node
 157     if node.type == 'GROUP_OUTPUT':
 158         iterator = node.id_data.interface.items_tree
 159     elif node.type == 'GROUP_INPUT':
 160         iterator = node.id_data.interface.items_tree
 161     elif hasattr(node, "node_tree"):
 162         iterator = node.node_tree.interface.items_tree
 163     else:
 164         return None
 165
 166     for s in iterator:
 167         if s.identifier == socket.identifier:
 168             return s
 169     return iterator[0]
 170
 171
 172 def is_viewer_link(link, output_node):
 173     if link.to_node == output_node and link.to_socket == output_node.inputs[0]:
 174         return True
 175     if link.to_node.type == 'GROUP_OUTPUT':
 176         socket = get_internal_socket(link.to_socket)
 177         if is_viewer_socket(socket):
 178             return True
 179     return False
 180
 181
 182 def get_group_output_node(tree, output_node_type='GROUP_OUTPUT'):
 183     for node in tree.nodes:
 184         if node.type == output_node_type and node.is_active_output:
 185             return node
 186
 187
 188 def get_output_location(tree):
 189     # get right-most location
 190     sorted_by_xloc = (sorted(tree.nodes, key=lambda x: x.location.x))
 191     max_xloc_node = sorted_by_xloc[-1]
 192
 193     # get average y location
 194     sum_yloc = 0
 195     for node in tree.nodes:
 196         sum_yloc += node.location.y
 197
 198     loc_x = max_xloc_node.location.x + max_xloc_node.dimensions.x + 80
 199     loc_y = sum_yloc / len(tree.nodes)
 200     return loc_x, loc_y
 201
 202
 203 def nw_check(cls, context):
 204     space = context.space_data
 205     if space.type != 'NODE_EDITOR':
 206         cls.poll_message_set("Current editor is not a node editor.")
 207         return False
 208     if space.node_tree is None:
 209         cls.poll_message_set("No node tree was found in the current node editor.")
 210         return False
 211     if space.node_tree.library is not None:
 212         cls.poll_message_set("Current node tree is linked from another .blend file.")
 213         return False
 214     return True
 215
 216
 217 def nw_check_not_empty(cls, context):
 218     if not context.space_data.node_tree.nodes:
 219         cls.poll_message_set("Current node tree does not contain any nodes.")
 220         return False
 221     return True
 222
 223
 224 def nw_check_active(cls, context):
 225     if context.active_node is None or not context.active_node.select:
 226         cls.poll_message_set("No active node.")
 227         return False
 228     return True
 229
 230
 231 def nw_check_selected(cls, context, min=1, max=inf):
 232     num_selected = len(context.selected_nodes)
 233     if num_selected < min:
 234         if min > 1:
 235             cls.poll_message_set(f"At least {min} nodes must be selected.")
 236         else:
 237             cls.poll_message_set(f"At least {min} node must be selected.")
 238         return False
 239     if num_selected > max:
 240         cls.poll_message_set(f"{num_selected} nodes are selected, but this operator can only work on {max}.")
 241         return False
 242     return True
 243
 244
 245 def nw_check_space_type(cls, context, types):
 246     if context.space_data.tree_type not in types:
 247         tree_types_str = ", ".join(t.split('NodeTree')[0].lower() for t in sorted(types))
 248         cls.poll_message_set("Current node tree type not supported.\n"
 249                              "Should be one of " + tree_types_str + ".")
 250         return False
 251     return True
 252
 253
 254 def nw_check_node_type(cls, context, type, invert=False):
 255     if invert and context.active_node.type == type:
 256         cls.poll_message_set(f"Active node should be not of type {type}.")
 257         return False
 258     elif not invert and context.active_node.type != type:
 259         cls.poll_message_set(f"Active node should be of type {type}.")
 260         return False
 261     return True
 262
 263
 264 def nw_check_visible_outputs(cls, context):
 265     if not any(is_visible_socket(out) for out in context.active_node.outputs):
 266         cls.poll_message_set("Current node has no visible outputs.")
 267         return False
 268     return True
 269
 270
 271 def nw_check_viewer_node(cls):
 272     for img in bpy.data.images:
 273         # False if not connected or connected but no image
 274         if (img.source == 'VIEWER'
 275                 and len(img.render_slots) == 0
 276                 and sum(img.size) > 0):
 277             return True
 278     cls.poll_message_set("Viewer image not found.")
 279     return False
 280
 281
 282 def get_first_enabled_output(node):
 283     for output in node.outputs:
 284         if output.enabled:
 285             return output
 286     else:
 287         return node.outputs[0]
 288
 289
 290 def is_visible_socket(socket):
 291     return not socket.hide and socket.enabled and socket.type != 'CUSTOM'
 292
 293
 294 class NWBase:
 295     @classmethod
 296     def poll(cls, context):
 297         return nw_check(cls, context)
 298
 299
 300 class NWBaseMenu:
 301     @classmethod
 302     def poll(cls, context):
 303         space = context.space_data
 304         return (space.type == 'NODE_EDITOR'
 305                 and space.node_tree is not None
 306                 and space.node_tree.library is None)