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
   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_active_tree(context):
 142     tree = context.space_data.node_tree
 143     path = []
 144     # Get nodes from currently edited tree.
 145     # If user is editing a group, space_data.node_tree is still the base level (outside group).
 146     # context.active_node is in the group though, so if space_data.node_tree.nodes.active is not
 147     # the same as context.active_node, the user is in a group.
 148     # Check recursively until we find the real active node_tree:
 149     if tree.nodes.active:
 150         while tree.nodes.active != context.active_node:
 151             tree = tree.nodes.active.node_tree
 152             path.append(tree)
 153     return tree, path
 154
 155
 156 def get_nodes_links(context):
 157     tree, path = get_active_tree(context)
 158     return tree.nodes, tree.links
 159
 160
 161 viewer_socket_name = "tmp_viewer"
 162
 163
 164 def is_viewer_socket(socket):
 165     # checks if a internal socket is a valid viewer socket
 166     return socket.name == viewer_socket_name and socket.NWViewerSocket
 167
 168
 169 def get_internal_socket(socket):
 170     # get the internal socket from a socket inside or outside the group
 171     node = socket.node
 172     if node.type == 'GROUP_OUTPUT':
 173         iterator = node.id_data.interface.items_tree
 174     elif node.type == 'GROUP_INPUT':
 175         iterator = node.id_data.interface.items_tree
 176     elif hasattr(node, "node_tree"):
 177         iterator = node.node_tree.interface.items_tree
 178     else:
 179         return None
 180
 181     for s in iterator:
 182         if s.identifier == socket.identifier:
 183             return s
 184     return iterator[0]
 185
 186
 187 def is_viewer_link(link, output_node):
 188     if link.to_node == output_node and link.to_socket == output_node.inputs[0]:
 189         return True
 190     if link.to_node.type == 'GROUP_OUTPUT':
 191         socket = get_internal_socket(link.to_socket)
 192         if is_viewer_socket(socket):
 193             return True
 194     return False
 195
 196
 197 def get_group_output_node(tree):
 198     for node in tree.nodes:
 199         if node.type == 'GROUP_OUTPUT' and node.is_active_output:
 200             return node
 201
 202
 203 def get_output_location(tree):
 204     # get right-most location
 205     sorted_by_xloc = (sorted(tree.nodes, key=lambda x: x.location.x))
 206     max_xloc_node = sorted_by_xloc[-1]
 207
 208     # get average y location
 209     sum_yloc = 0
 210     for node in tree.nodes:
 211         sum_yloc += node.location.y
 212
 213     loc_x = max_xloc_node.location.x + max_xloc_node.dimensions.x + 80
 214     loc_y = sum_yloc / len(tree.nodes)
 215     return loc_x, loc_y
 216
 217
 218 def nw_check(context):
 219     space = context.space_data
 220     valid_trees = ["ShaderNodeTree", "CompositorNodeTree", "TextureNodeTree", "GeometryNodeTree"]
 221
 222     if (space.type == 'NODE_EDITOR'
 223             and space.node_tree is not None
 224             and space.node_tree.library is None
 225             and space.tree_type in valid_trees):
 226         return True
 227
 228     return False
 229
 230
 231 def get_first_enabled_output(node):
 232     for output in node.outputs:
 233         if output.enabled:
 234             return output
 235     else:
 236         return node.outputs[0]
 237
 238
 239 def is_visible_socket(socket):
 240     return not socket.hide and socket.enabled and socket.type != 'CUSTOM'
 241
 242
 243 class NWBase:
 244     @classmethod
 245     def poll(cls, context):
 246         return nw_check(context)