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 f34cbe1566e35168f36e77845cccd160a075681b..5731671babccb5b3850258443b8eceb57451e569 100644 (file)
--- 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
 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'],
 
     node_tree.links.new(
         node_tree.nodes['From Node'].outputs['Output'],


@@ -15,9 +15,20 @@ something like::
 
 Now this can look 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'] \
 
     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')
 
 
     remove_links(node_tree, '+Node')
 

-Outputs::
+Or just outputs::

 
     remove_links(node_tree, '-Node')
 
 
     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]")
 
 
     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.
 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
 versions.
\ No newline at end of file

diff --git a/blender_node_tools.py b/blender_node_tools.py
index 4be6a9c8c20ccc2d6ecccd13266750d33af4500c..03be748d67128eefbd5d0c9c7ee18ec5ca5ca075 100644 (file)
--- a/blender_node_tools.py
+++ b/blender_node_tools.py
@@ -4,6 +4,7 @@
 import sys
 
 from typing import (
 import sys
 
 from typing import (

+    Dict,

     Tuple,
     List,
     Optional,
     Tuple,
     List,
     Optional,


@@ -51,27 +52,50 @@ NodeType = Union[str, Node, NodeSocket]
 
 INPUT_SYMBOL = '+'
 OUTPUT_SYMBOL = '-'
 
 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,
 
 def link_nodes(
     tree: NodeTree,


@@ -91,32 +115,11 @@ def link_nodes(
     """
     connections = 0
 
     """
     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
 
     for output_data in output_sockets:
         output, out_matched = output_data


@@ -136,86 +139,17 @@ def link_nodes(
 
     return connections
 
 
     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
     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 ↓
 
 
 # helpers ↓
 


@@ -224,22 +158,17 @@ def _append_matching_sockets(
     socket_index: str,
     include: bool,
     node_sockets: Union[NodeInputs,NodeOutputs],
     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)
     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
 
                 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))
 
 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:]
 
         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]
     if len(parts) > 1:
         socket = parts[-1]

-        node = SOCKET_DELIMETER.join(parts[:-1])
+        node = SOCKET_DELIMITER.join(parts[:-1])

     else:
         node = node_str
 
     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]
 
             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,
+    }