Merge pull request #152 from Hoikas/nodes

Node Improvements

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,52 @@ class PlasmaNodeBase:
                     if idname == node.bl_idname:
                         yield i
 
+    def generate_valid_links_for(self, context, socket, is_output):
+        """Generates valid node sockets that can be linked to a specific socket on this node."""
+        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
+
+            # Korman standard node socket definitions
+            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"] }
+
+            # Some node types (eg Python) may auto-generate their own sockets, so we ask them now.
+            for i in dest_node_cls.generate_valid_links_to(context, socket, is_output):
+                yield i
+
+    @classmethod
+    def generate_valid_links_to(cls, context, socket, is_output):
+        """Generates valid sockets on this node type that can be linked to a specific node's socket."""
+        return []
+
     def harvest_actors(self, bo):
         return set()
 
@@ -330,7 +376,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.generate_valid_links_for(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 +397,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,8 +21,9 @@ 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
+from ..plasma_attributes import get_attributes_from_str
 
 _single_user_attribs = {
     "ptAttribBoolean", "ptAttribInt", "ptAttribFloat", "ptAttribString", "ptAttribDropDownList",
@@ -306,6 +307,88 @@ class PlasmaPythonFileNode(PlasmaVersionedNode, bpy.types.Node):
             if i.attribute_id == idx:
                 yield i
 
+    def generate_valid_links_for(self, context, socket, is_output):
+        # Python nodes have no outputs...
+        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"] }
+
+    @classmethod
+    def generate_valid_links_to(cls, context, socket, is_output):
+        # This is only useful for nodes wanting to connect to our inputs (ptAttributes)
+        if not is_output:
+            return
+
+        if isinstance(socket, PlasmaPythonAttribNodeSocket):
+            pl_attrib = socket.node.pl_attrib
+        else:
+            pl_attrib = getattr(socket.node, "pl_attrib", set())
+            if not pl_attrib:
+                return
+
+            # Fetch the output definition for the requested socket and make sure it can connect to us.
+            socket_def = getattr(socket.node, "output_sockets", {}).get(socket.alias)
+            if socket_def is None:
+                return
+            valid_link_sockets = socket_def.get("valid_link_sockets")
+            valid_link_nodes = socket_def.get("valid_link_nodes")
+            if valid_link_sockets is not None and "PlasmaPythonFileNodeSocket" not in valid_link_sockets:
+                return
+            if valid_link_nodes is not None and "PlasmaPythonFileNode" not in valid_link_nodes:
+                return
+
+        # Ok, apparently this thing can connect as a ptAttribute. The only problem with that is
+        # that we have no freaking where... The sockets are spawned by Python files... So, we
+        # need to look at all the Python files we know about...
+        for text_id in bpy.data.texts:
+            if not text_id.name.endswith(".py"):
+                continue
+            attribs = get_attributes_from_str(text_id.as_string())
+            if not attribs:
+                continue
+
+            for _, attrib in attribs.items():
+                if not attrib["type"] in pl_attrib:
+                    continue
+
+                # *gulp*
+                yield { "node_idname": "PlasmaPythonFileNode",
+                        "node_text": text_id.name,
+                        "node_settings": { "filename": text_id.name },
+                        "socket_name":  attrib["name"],
+                        "socket_text": attrib["name"] }
+
     def harvest_actors(self, bo):
         actors = set()
         actors.add(bo.name)
@@ -421,7 +504,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 +525,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)

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

@@ -16,6 +16,7 @@
 import bpy
 from bpy.props import *
 import itertools
+import pickle
 
 class NodeOperator:
     @classmethod
@@ -23,6 +24,119 @@ 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.generate_valid_links_for(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):
+            # Pickle protocol 0 uses only ASCII bytes, so we can pretend it's a string easily...
+            id_string = pickle.dumps(link, protocol=0).decode()
+            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):
+        link = pickle.loads(node_item.encode())
+        self._hack.clear()
+
+        tree = context.space_data.edit_tree
+        dest_node = tree.nodes.new(type=link["node_idname"])
+        for attr, value in link.get("node_settings", {}).items():
+            setattr(dest_node, attr, value)
+        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(link["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"