Add node create+link operator.

Inspired by Unreal Blueprint's drag link+create list functionality. I
realize Blender has a few operators that do different parts of this job,
but they don't provide this well-polished functionality.

korman/nodes/node_conditions.py

@@ -258,7 +258,7 @@ class PlasmaFacingTargetSocket(PlasmaNodeSocketBase, bpy.types.NodeSocket):
                            description="Avatar must be facing the target object",
                            default=True)
 
-    def draw(self, context, layout, node, text):
+    def draw_content(self, context, layout, node, text):
         if self.simple_mode:
             layout.prop(self, "allow_simple", text="")
         layout.label(text)
@@ -505,7 +505,7 @@ class PlasmaVolumeSettingsSocket(PlasmaNodeSocketBase):
 class PlasmaVolumeSettingsSocketIn(PlasmaVolumeSettingsSocket, bpy.types.NodeSocket):
     allow = BoolProperty()
 
-    def draw(self, context, layout, node, text):
+    def draw_content(self, context, layout, node, text):
         if not self.is_linked:
             layout.prop(self, "allow", text="")
         layout.label(text)

korman/nodes/node_core.py

@@ -139,6 +139,41 @@ class PlasmaNodeBase:
                     if idname == node.bl_idname:
                         yield i
 
+    def get_valid_link_search(self, context, socket, is_output):
+        from .node_deprecated import PlasmaDeprecatedNode
+
+        for dest_node_cls in bpy.types.Node.__subclasses__():
+            if not issubclass(dest_node_cls, PlasmaNodeBase) or issubclass(dest_node_cls, PlasmaDeprecatedNode):
+                continue
+            socket_defs = getattr(dest_node_cls, "input_sockets", {}) if is_output else \
+                          getattr(dest_node_cls, "output_sockets", {})
+
+            for socket_name, socket_def in socket_defs.items():
+                if socket_def.get("can_link") is False:
+                    continue
+                if socket_def.get("hidden") is True:
+                    continue
+                
+                valid_source_nodes = socket_def.get("valid_link_nodes")
+                valid_source_sockets = socket_def.get("valid_link_sockets")
+                if valid_source_nodes is not None and self.bl_idname not in valid_source_nodes:
+                    continue
+                if valid_source_sockets is not None and socket.bl_idname not in valid_source_sockets:
+                    continue
+                if valid_source_sockets is None and valid_source_nodes is None:
+                    if socket.bl_idname != socket_def["type"]:
+                        continue
+
+                # Can we even add the node?
+                poll_add = getattr(dest_node_cls, "poll_add", None)
+                if poll_add is not None and not poll_add(context):
+                    continue
+
+                yield { "node_idname": dest_node_cls.bl_idname,
+                        "node_text": dest_node_cls.bl_label,
+                        "socket_name": socket_name,
+                        "socket_text": socket_def["text"] }
+
     def harvest_actors(self, bo):
         return set()
 
@@ -330,7 +365,20 @@ class PlasmaNodeSocketBase:
         return ident.rsplit('.', 1)[0]
 
     def draw(self, context, layout, node, text):
-        layout.label(text)
+        if not self.is_output:
+            self.draw_add_operator(context, layout, node)
+        self.draw_content(context, layout, node, text)
+        if self.is_output:
+            self.draw_add_operator(context, layout, node)
+
+    def draw_add_operator(self, context, layout, node):
+        row = layout.row()
+        row.enabled = any(node.get_valid_link_search(context, self, self.is_output))
+        row.operator_context = "INVOKE_DEFAULT"
+        add_op = row.operator("node.plasma_create_link_node", text="", icon="ZOOMIN")
+        add_op.node_name = node.name
+        add_op.sock_ident = self.identifier
+        add_op.is_output = self.is_output
 
     def draw_color(self, context, node):
         # It's so tempting to just do RGB sometimes... Let's be nice.
@@ -338,6 +386,9 @@ class PlasmaNodeSocketBase:
             return tuple(self.bl_color[0], self.bl_color[1], self.bl_color[2], 1.0)
         return self.bl_color
 
+    def draw_content(self, context, layout, node, text):
+        layout.label(text)
+
     @property
     def is_used(self):
         return bool(self.links)

korman/nodes/node_python.py

@@ -21,7 +21,7 @@ from PyHSPlasma import *
 
 from ..korlib import replace_python2_identifier
 from .node_core import *
-from .node_deprecated import PlasmaVersionedNode
+from .node_deprecated import PlasmaDeprecatedNode, PlasmaVersionedNode
 from .. import idprops
 
 _single_user_attribs = {
@@ -306,6 +306,42 @@ class PlasmaPythonFileNode(PlasmaVersionedNode, bpy.types.Node):
             if i.attribute_id == idx:
                 yield i
 
+    def get_valid_link_search(self, context, socket, is_output):
+        assert is_output is False
+
+        attrib_type = socket.attribute_type
+        for i in bpy.types.Node.__subclasses__():
+            node_attrib_types = getattr(i, "pl_attrib", None)
+            if node_attrib_types is None or issubclass(i, PlasmaDeprecatedNode):
+                continue
+
+            if attrib_type in node_attrib_types:
+                if issubclass(i, PlasmaAttribNodeBase):
+                   yield { "node_idname": i.bl_idname,
+                           "node_text": i.bl_label,
+                           "socket_name": "pfm",
+                           "socket_text": "Python File" }
+                else:
+                    for socket_name, socket_def in i.output_sockets.items():
+                        if socket_def.get("hidden") is True:
+                            continue
+                        if socket_def.get("can_link") is False:
+                            continue
+
+                        valid_link_nodes = socket_def.get("valid_link_nodes")
+                        valid_link_sockets = socket_def.get("valid_link_sockets")
+                        if valid_link_nodes is not None and self.bl_idname not in valid_link_nodes:
+                            print(socket_name, self.bl_idname, valid_link_nodes)
+                            continue
+                        if valid_link_sockets is not None and "PlasmaPythonFileNodeSocket" not in valid_link_sockets:
+                            print(socket_name, "PlasmaPythonFileNodeSocket", valid_link_sockets)
+                            continue
+
+                        yield { "node_idname": i.bl_idname,
+                                "node_text": i.bl_label,
+                                "socket_name": socket_name,
+                                "socket_text": socket_def["text"] }
+
     def harvest_actors(self, bo):
         actors = set()
         actors.add(bo.name)
@@ -421,7 +457,7 @@ class PlasmaPythonFileNodeSocket(PlasmaNodeSocketBase, bpy.types.NodeSocket):
         return self.node.attribute_map[self.attribute_id].attribute_type
 
     def draw(self, context, layout, node, text):
-        layout.alignment = "LEFT"
+        self.draw_add_operator(context, layout, node)
         layout.label("ID: {}".format(self.attribute_id))
         layout.label(self.attribute_description)
 
@@ -442,7 +478,7 @@ class PlasmaPythonFileNodeSocket(PlasmaNodeSocketBase, bpy.types.NodeSocket):
 
 
 class PlasmaPythonAttribNodeSocket(PlasmaNodeSocketBase, bpy.types.NodeSocket):
-    def draw(self, context, layout, node, text):
+    def draw_content(self, context, layout, node, text):
         attrib = node.to_socket
         if attrib is None:
             layout.label(text)
@@ -466,6 +502,10 @@ class PlasmaAttribNodeBase(PlasmaNodeBase):
         attr = self.to_socket
         return "Value" if attr is None else attr.attribute_name
 
+    def get_valid_link_search(self, context, socket, is_output):
+        # This quick'n'dirty hack disables the + button on all sockets.
+        return []
+
     @property
     def to_socket(self):
         """Returns the socket linked to IF only one link has been made"""

korman/nodes/node_responder.py

@@ -355,7 +355,7 @@ class PlasmaRespStateSocket(PlasmaNodeSocketBase, bpy.types.NodeSocket):
 class PlasmaRespStateRefSocket(PlasmaNodeSocketBase, bpy.types.NodeSocket):
     bl_color = (1.00, 0.980, 0.322, 1.0)
 
-    def draw(self, context, layout, node, text):
+    def draw_content(self, context, layout, node, text):
         if isinstance(node, PlasmaResponderNode):
             try:
                 idx = next((idx for idx, socket in enumerate(node.find_output_sockets("state_refs")) if socket == self))

korman/operators/op_nodes.py

@@ -23,6 +23,116 @@ class NodeOperator:
         return context.scene.render.engine == "PLASMA_GAME"
 
 
+class CreateLinkNodeOperator(NodeOperator, bpy.types.Operator):
+    bl_idname = "node.plasma_create_link_node"
+    bl_label = "Create Node"
+    bl_description = "Create and link a new node to this socket"
+    bl_options = {"UNDO", "INTERNAL"}
+    bl_property = "node_item"
+
+    node_name = StringProperty()
+    sock_ident = StringProperty()
+    is_output = BoolProperty()
+
+    # The "official" node search operator does something like this...
+    # Documentation seems to indicate this works around poor refcounting.
+    _hack = []
+
+    def _link_search_list(self, context):
+        CreateLinkNodeOperator._hack = list(CreateLinkNodeOperator._link_search_list_imp(self, context))
+        return CreateLinkNodeOperator._hack
+
+    def _link_search_list_imp(self, context):
+        # NOTE: `self` is not actually an instance of this class. It's a fancy wrapper object
+        # whose only members are the above properties...
+        tree = context.space_data.edit_tree
+        src_node = tree.nodes[self.node_name]
+        src_socket = CreateLinkNodeOperator._find_source_socket(self, src_node)
+
+        links = list(src_node.get_valid_link_search(context, src_socket, self.is_output))
+        max_node = max((len(i["node_text"]) for i in links)) if links else 0
+        for i, link in enumerate(links):
+            id_string = "{}!@!{}".format(link["node_idname"], link["socket_name"])
+            desc_string = "{node}:{node_sock_space}{sock}".format(node=link["node_text"],
+                node_sock_space=(" " * (max_node - len(link["node_text"]) + 4)),
+                sock=link["socket_text"])
+            yield (id_string, desc_string, "", i)
+
+    node_item = EnumProperty(items=_link_search_list)
+
+    def _find_source_socket(self, node):
+        sockets = node.outputs if self.is_output else node.inputs
+        for i in sockets:
+            if i.identifier == self.sock_ident:
+                return i
+        raise LookupError()
+
+    def invoke(self, context, event):
+        possible_links = self._link_search_list(context)
+        if not possible_links:
+            self.report({"WARNING"}, "No nodes can be created.")
+            return {"FINISHED"}
+        elif len(possible_links) == 1:
+            context.window_manager.modal_handler_add(self)
+            return {"RUNNING_MODAL"}
+        else:
+            context.window_manager.invoke_search_popup(self)
+            return {"RUNNING_MODAL"}
+
+    def execute(self, context):
+        context.window_manager.modal_handler_add(self)
+        return {"RUNNING_MODAL"}
+
+    def _create_link_node(self, context, node_item):
+        node_type, socket_name = node_item.split("!@!")
+        self._hack.clear()
+
+        tree = context.space_data.edit_tree
+        dest_node = tree.nodes.new(type=node_type)
+        for i in tree.nodes:
+            i.select = i == dest_node
+        tree.nodes.active = dest_node
+        dest_node.location = context.space_data.cursor_location
+
+        src_node = tree.nodes[self.node_name]
+        src_socket = self._find_source_socket(src_node)
+        # We need to use Korman's functions because they may generate a node socket.
+        find_socket = dest_node.find_input_socket if self.is_output else dest_node.find_output_socket
+        dest_socket = find_socket(socket_name, True)
+
+        if self.is_output:
+            tree.links.new(src_socket, dest_socket)
+        else:
+            tree.links.new(dest_socket, src_socket)
+        self.finished = True
+        return {"FINISHED"}
+
+    def modal(self, context, event):
+        # Ugh. The Blender API sucks so much. We can only get the cursor pos from here???
+        context.space_data.cursor_location_from_region(event.mouse_region_x, event.mouse_region_y)
+        if len(self._hack) == 1:
+            self._create_link_node(context, self._hack[0][0])
+            self._hack.clear()
+        elif self._hack:
+            self._create_link_node(context, self.node_item)
+            self._hack.clear()
+
+        if event.type == "MOUSEMOVE":
+            tree = context.space_data.edit_tree
+            tree.nodes.active.location = context.space_data.cursor_location
+        elif event.type in {"ESC", "LEFTMOUSE"}:
+            return {"FINISHED"}
+        return {"RUNNING_MODAL"}
+
+    @classmethod
+    def poll(cls, context):
+        space = context.space_data
+        # needs active node editor and a tree to add nodes to
+        return (space.type == 'NODE_EDITOR' and
+                space.edit_tree and not space.edit_tree.library and
+                context.scene.render.engine == "PLASMA_GAME")
+
+
 class SelectFileOperator(NodeOperator, bpy.types.Operator):
     bl_idname = "file.plasma_file_picker"
     bl_label = "Select"