diff --git a/korman/nodes/node_core.py b/korman/nodes/node_core.py index 1d48bda..e84a1fb 100644 --- a/korman/nodes/node_core.py +++ b/korman/nodes/node_core.py @@ -139,15 +139,17 @@ class PlasmaNodeBase: if idname == node.bl_idname: yield i - def get_valid_link_search(self, context, socket, is_output): + 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 @@ -174,6 +176,15 @@ class PlasmaNodeBase: "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() @@ -373,7 +384,7 @@ class PlasmaNodeSocketBase: 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.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 diff --git a/korman/nodes/node_python.py b/korman/nodes/node_python.py index 9393694..0780c6b 100644 --- a/korman/nodes/node_python.py +++ b/korman/nodes/node_python.py @@ -23,6 +23,7 @@ from ..korlib import replace_python2_identifier from .node_core import * 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,7 +307,8 @@ class PlasmaPythonFileNode(PlasmaVersionedNode, bpy.types.Node): if i.attribute_id == idx: yield i - def get_valid_link_search(self, context, socket, is_output): + 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 @@ -342,6 +344,51 @@ class PlasmaPythonFileNode(PlasmaVersionedNode, bpy.types.Node): "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) @@ -502,10 +549,6 @@ 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""" diff --git a/korman/operators/op_nodes.py b/korman/operators/op_nodes.py index ae8c753..67e1196 100644 --- a/korman/operators/op_nodes.py +++ b/korman/operators/op_nodes.py @@ -16,6 +16,7 @@ import bpy from bpy.props import * import itertools +import pickle class NodeOperator: @classmethod @@ -49,10 +50,11 @@ class CreateLinkNodeOperator(NodeOperator, bpy.types.Operator): 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)) + 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): - id_string = "{}!@!{}".format(link["node_idname"], link["socket_name"]) + # 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"]) @@ -84,11 +86,13 @@ class CreateLinkNodeOperator(NodeOperator, bpy.types.Operator): return {"RUNNING_MODAL"} def _create_link_node(self, context, node_item): - node_type, socket_name = node_item.split("!@!") + link = pickle.loads(node_item.encode()) self._hack.clear() tree = context.space_data.edit_tree - dest_node = tree.nodes.new(type=node_type) + 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 @@ -98,7 +102,7 @@ class CreateLinkNodeOperator(NodeOperator, bpy.types.Operator): 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) + dest_socket = find_socket(link["socket_name"], True) if self.is_output: tree.links.new(src_socket, dest_socket)