From: Dustin Walde Date: Tue, 17 May 2022 23:52:35 +0000 (-0700) Subject: Clean up API and Documentation (#1) X-Git-Url: https://git.walde.dev/?a=commitdiff_plain;h=5f7ae20a346425c345ae9563235ef0469de4e81e;p=blender_node_tools Clean up API and Documentation (#1) - Make API more consitent and concise - Update README to reflect changes - Clean up source file to be consisten with readme --- diff --git a/README.rst b/README.rst index f34cbe1..5731671 100644 --- a/README.rst +++ b/README.rst @@ -5,8 +5,8 @@ A python function module to aid in manipulating node connections in Blender. The goal is to make it simpler when adding / removing links between -nodes of a NodeTree. Using the Blender Python API, adding links looks -something like:: +nodes of a NodeTree. Using the Blender Python API as it exists, +adding links looks something like:: node_tree.links.new( node_tree.nodes['From Node'].outputs['Output'], @@ -15,9 +15,20 @@ something like:: Now this can look like:: - link_nodes(node_tree, 'From Node', 'To Node:Input') + link_nodes(node_tree, 'From Node:Output', 'To Node:Input') -Even more roundabout was the removal of links:: +Even simpler if the input and outputs have only one matching type::: + + link_nodes(node_tree, 'From Node', 'To Node') + +When linking nodes, if a node string resolves to multiple sockets, +``link_nodes`` will attempt to link each of them, in order until it +matches with one of the input sockets with the same type. +The only time sockets of different types are connected is when both +input and output resolve to a single socket. + +As another example, and even more roundabout in Blender +is the removal of links:: for link in node_tree.links: if link.to_node == node_tree.nodes['Node'] \ @@ -32,20 +43,76 @@ Or to remove all inputs:: remove_links(node_tree, '+Node') -Outputs:: +Or just outputs:: remove_links(node_tree, '-Node') +Syntax +------ + +``blender_node_tools`` uses a custom string format as a shorthand +to access ``Node``s and ``NodeSocket``s in a Blender ``NodeTree``. + +At a baseline, a string represents a Node name.:: + + "DemoNode" + +On its own, this refers to all sockets of the node. +Depending on the context, input or output sockets can be inferred. +To specify only inputs or only outputs, a leading ``+`` (Inputs) +or `-` (Outputs) can be used.:: + + "+DemoNode" + +Finally, to refer to a specific socket, the socket name can be placed +after the node name delimited by a ``:``:: + + "+DemoNode:DemoSocket" + + "-DemoNode:DemoSocket" + +And if the socket name is unique to either side, the leading +input/output specifier is unneeded:: + + "DemoNode:DemoOutputSocket" -Something like a compositor mix node have two inputs of the same name -"Image". These may be targeting with array index notation:: +Something like a compositor mix node may have two inputs of the same +name (e.g. "Image"). +These may be targeted with array index notation:: link_nodes(node_tree, "Image Texture", "Mix:Image[0]") link_nodes(node_tree, "Fade Texture", "Mix:Image[1]") +Functions +--------- + +``NodeType = Union[str, Node, NodeSocket]`` + +Data Access: + +``get_node(tree: NodeTree, of: NodeType) -> Node`` + +``get_sockets(tree: NodeTree, node: NodeType) -> List[NodeSocket]`` + +``get_links(tree: NodeTree, node: NodeType) -> List[NodeLink]`` + +There are also singular ``get_socket`` and ``get_link`` functions. + +``get_subtree(tree: NodeTree, node: NodeType) -> Optional[NodeTree]`` + +``is_group(node: Node) -> bool`` + +Graph modification: + +``link_nodes(tree: NodeTree, node_from: NodeType, node_to: NodeType, preserve_existing: bool = False) -> int`` + +``remove_links(tree: NodeTree, node: NodeType, node2: Optional[NodeType] = None) -> int`` + +If you already have the exact ``NodeLink``, use ``node_tree.links.remove(link)``. + Requirements ------------ Python version 3.6 or later. -Tested with Blender version 2.92, will likely work for previous +Tested with Blender version 2.92, 2.93, and 3.1, will likely work for previous versions. \ No newline at end of file diff --git a/blender_node_tools.py b/blender_node_tools.py index 4be6a9c..03be748 100644 --- a/blender_node_tools.py +++ b/blender_node_tools.py @@ -4,6 +4,7 @@ import sys from typing import ( + Dict, Tuple, List, Optional, @@ -51,27 +52,50 @@ NodeType = Union[str, Node, NodeSocket] INPUT_SYMBOL = '+' OUTPUT_SYMBOL = '-' -SOCKET_DELIMETER = ':' +SOCKET_DELIMITER = ':' -def is_group(node: Node) -> bool: +def get_node(tree: NodeTree, of: NodeType) -> Node: """ - Check if node is a *GroupNode of some sort. + Retrieve the the Node from tree at the given `node_str`. + + raises ValueError if `node_path` is invalid. """ - return _is_union_instance(node, NodeGroupType) + return _resolve_node(tree, of)['node'] + +def get_sockets(tree: NodeTree, of: str) -> List[NodeSocket]: + node_data = _resolve_node(tree, of) + node_data['input_sockets'].extend(node_data['output_sockets']) + return node_data['input_sockets'] + +def get_socket(tree: NodeTree, of: str) -> Optional[NodeSocket]: + sockets = get_sockets(tree, of) + if len(sockets) == 0: + return None + return sockets[0] -def format_node(node: str, - input: Union[bool,str]=False, - output: Union[bool,str]=False) \ - -> str: - if input != False and output != False: - raise ValueError("Node ref cannot be both input and output") +def get_links(tree: NodeTree, of: NodeType) -> List[NodeLink]: + node_data = _resolve_node(tree, of) + node_data['input_links'].extend(node_data['output_links']) + return node_data['input_links'] - if input != False: - return _format_socket_str(node, input, INPUT_SYMBOL) - elif output != False: - return _format_socket_str(node, output, OUTPUT_SYMBOL) +def get_link(tree: NodeTree, of: NodeType) -> Optional[NodeLink]: + links = get_links(tree, of) + if len(links) == 0: + return None + return links[0] + +def get_subtree(tree: NodeTree, of: NodeType) -> Optional[NodeTree]: + node = _resolve_node(tree, of)['node'] + if is_group(node): + return node.node_tree + else: + return None - return node +def is_group(node: Node) -> bool: + """ + Check if node is a *GroupNode of some sort. + """ + return _is_union_instance(node, NodeGroupType) def link_nodes( tree: NodeTree, @@ -91,32 +115,11 @@ def link_nodes( """ connections = 0 - output_sockets = [] - input_sockets = [] - - # get output sockets - if isinstance(link_from, NodeSocket): - if link_from.is_output: - output_sockets.append(link_from) - elif isinstance(link_from, Node): - for output in link_from.outputs: - output_sockets.append(output) - else: - _, _, output_sockets = resolve_node(tree, link_from) - - output_sockets = _label_sockets(output_sockets, preserve_existing) - - # get input sockets - if isinstance(link_to, NodeSocket): - if not link_to.is_output: - input_sockets.append(link_to) - elif isinstance(link_to, Node): - for input in link_to.inputs: - input_sockets.append(input) - else: - _, input_sockets, _ = resolve_node(tree, link_to) + from_data = _resolve_node(tree, link_from) + output_sockets = _label_sockets(from_data['output_sockets'], preserve_existing) - input_sockets = _label_sockets(input_sockets, preserve_existing) + to_data = _resolve_node(tree, link_to) + input_sockets = _label_sockets(to_data['input_sockets'], preserve_existing) for output_data in output_sockets: output, out_matched = output_data @@ -136,86 +139,17 @@ def link_nodes( return connections -def get_link(tree: NodeTree, of: NodeType) -> Optional[NodeLink]: - links = get_links(tree, of) - if len(links) == 0: - return None - return links[0] - -def get_links(tree: NodeTree, of: NodeType) -> List[NodeLink]: - links = [] - if isinstance(of, NodeSocket): - links = of.links - elif isinstance(of, Node): - for input in of.inputs: - links.extend(input.links) - for output in of.outputs: - links.extend(output.links) - else: - _, inputs, outputs = resolve_node(tree, of) - for input in inputs: - links.extend(input.links) - for output in outputs: - links.extend(output.links) - - return links - -def remove_links(tree: NodeTree, node: NodeType) -> int: +def remove_links(tree: NodeTree, of: NodeType, to: Optional[NodeType] = None) -> int: removed = 0 - for link in get_links(tree, node): - tree.links.remove(link) - removed += 1 - - return removed - -def get_socket(tree: NodeTree, node_str: str) -> NodeSocket: - return get_sockets(tree, node_str)[0] - -def get_sockets(tree: NodeTree, node_str: str) -> List[NodeSocket]: - _, inputs, outputs = resolve_node(tree, node_str) - inputs.extend(outputs) - return inputs - -def get_input_sockets(tree: NodeTree, node_str: str) -> List[NodeSocket]: - return resolve_node(tree, node_str)[1] - -def get_output_sockets(tree: NodeTree, node_str: str) -> List[NodeSocket]: - return resolve_node(tree, node_str)[2] - -def get_node(tree: NodeTree, node_str: str) -> Node: - """ - Retrieve the the Node from tree at the given `node_str`. - - raises ValueError if `node_path` is invalid. - """ - return resolve_node(tree, node_str)[0] - -def resolve_node(tree: NodeTree, node_str: str) \ - -> Tuple[Node, List[NodeSocket], List[NodeSocket]]: - """ - In the context of `tree`, get the Node and input and/or output sockets - of the `node_str`. - - Return: - Tuple containing (Node, `Socket`s, output `Socket`s) - """ - - node_name, socket_name, socket_index, include_inputs, include_outputs = \ - _parse_node_str(node_str) - - if node_name not in tree.nodes: - _invalid_node_str(node_str, - "Node name ({}) not found".format(node_name)) + if to is not None: + to = _resolve_node(tree, to)["node"] - node = tree.nodes[node_name] + for link in get_links(tree, of): + if to is None or link.to_node == to: + tree.links.remove(link) + removed += 1 - inputs = [] - outputs = [] - - _append_matching_sockets(socket_name, socket_index, include_inputs, node.inputs, inputs) - _append_matching_sockets(socket_name, socket_index, include_outputs, node.outputs, outputs) - - return (node, inputs, outputs) + return removed # helpers ↓ @@ -224,22 +158,17 @@ def _append_matching_sockets( socket_index: str, include: bool, node_sockets: Union[NodeInputs,NodeOutputs], - out_sockets: List[NodeSocket]): + out_sockets: List[NodeSocket], + out_links: List[NodeLink]): name_count = 0 if include: for socket in node_sockets: if socket_name is None or socket_name == socket.name: if socket_index < 0 or socket_index == name_count: out_sockets.append(socket) + out_links.extend(socket.links) name_count += 1 -def _format_socket_str(node, socket, symbol): - if isinstance(socket, str): - return '{}{}{}{}'.format( - symbol, node, SOCKET_DELIMETER, socket) - - return '{}{}'.format(symbol, node) - def _invalid_node_str(node_str, message): raise ValueError("Invalid Node Str: {}\n{}".format(node_str, message)) @@ -276,10 +205,10 @@ def _parse_node_str(node_str: str) -> Tuple[str, str, int, bool, bool]: input = False node_str = node_str[1:] - parts = node_str.split(SOCKET_DELIMETER) + parts = node_str.split(SOCKET_DELIMITER) if len(parts) > 1: socket = parts[-1] - node = SOCKET_DELIMETER.join(parts[:-1]) + node = SOCKET_DELIMITER.join(parts[:-1]) else: node = node_str @@ -291,4 +220,52 @@ def _parse_node_str(node_str: str) -> Tuple[str, str, int, bool, bool]: socket_index = int(socket[left_bracket+1:right_bracket]) socket = socket[:left_bracket] - return (node, socket, socket_index, input, output) \ No newline at end of file + return (node, socket, socket_index, input, output) + +def _resolve_node(tree: NodeTree, node: NodeType) \ + -> Dict[str, Union[Node, List[NodeSocket], List[NodeLink]]]: + """ + In the context of ``tree``, get the Node and input and/or output sockets + of the `node`. + + Return: + Tuple containing (Node, ``Socket``s, output ``Socket``s) + """ + + inputs = [] + outputs = [] + input_links = [] + output_links = [] + + if isinstance(node, str): + node_name, socket_name, socket_index, include_inputs, include_outputs = \ + _parse_node_str(node) + + if node_name not in tree.nodes: + _invalid_node_str(node, + "Node name ({}) not found".format(node_name)) + + node = tree.nodes[node_name] + _append_matching_sockets(socket_name, socket_index, include_inputs, node.inputs, inputs, input_links) + _append_matching_sockets(socket_name, socket_index, include_outputs, node.outputs, outputs, output_links) + elif isinstance(node, Node): + outputs.extend(node.outputs) + inputs.extend(node.inputs) + input_links.extend(filter(lambda a: a.to_node == node, node.links)) + output_links.extend(filter(lambda a: a.from_node == node, node.links)) + elif isinstance(node, NodeSocket): + if node.is_output: + outputs.append(node) + output_links.extend(node.links) + else: + inputs.append(node) + input_links.extend(node.links) + node = node.node + + return { + "node": node, + "input_links": input_links, + "input_sockets": inputs, + "output_links": output_links, + "output_sockets": outputs, + }