nodeslicer/infra/blender/list_utils.py

623 lines
24 KiB
Python

from __future__ import annotations
from typing import Any, Callable
import numpy as np
LIST_CAPABLE_SOCKET_TYPES = {
"SLICESOCK_Mesh",
"SLICESOCK_Paths",
"SLICESOCK_Object",
"SLICESOCK_Geometry",
"SLICESOCK_Intrinsics",
"SLICESOCK_Settings",
"NodeSocketObject",
"NodeSocketString",
"NodeSocketVector",
"NodeSocketRotation",
"NodeSocketFloat",
"NodeSocketInt",
"NodeSocketColor",
"NodeSocketBool",
}
def is_vector_value(val: Any) -> bool:
"""Checks if val represents a single 3D/4D vector, Euler rotation, or color,
rather than a list of such values."""
if val is None:
return False
val_type = type(val).__name__
if val_type in ("Vector", "Euler", "Quaternion", "Color"):
return True
if isinstance(val, np.ndarray):
if val.ndim == 1 and val.shape[0] in (3, 4):
return True
return False
if isinstance(val, (list, tuple)):
if len(val) in (3, 4) and all(isinstance(x, (int, float, np.number)) and not isinstance(x, bool) for x in val):
return True
return False
def is_input_a_list(val: Any, is_vector: bool = False) -> bool:
"""Determines whether an input value represents a list of items.
If is_vector is True, a single vector/rotation/color is treated as an atomic single value."""
if val is None or isinstance(val, (str, bytes, bool)):
return False
if type(val).__name__ == "SlicePaths" or getattr(val, "is_slice_paths", False):
return False
if is_vector:
if is_vector_value(val):
return False
return isinstance(val, (list, tuple))
return isinstance(val, (list, tuple))
def is_list_data(val: Any) -> bool:
"""Checks if val represents a list of items (Object, Mesh, Path, Gcode).
Note: SlicePaths is a list subclass representing an individual toolpath collection,
not a list of toolpaths, so it is treated as an atomic item (not list data)."""
return is_input_a_list(val, is_vector=False)
def flatten_list(items: Any) -> list[Any]:
"""Recursively flattens nested lists/tuples, keeping atomic items intact."""
result: list[Any] = []
if not is_list_data(items):
if items is not None:
result.append(items)
return result
for item in items:
if is_list_data(item):
result.extend(flatten_list(item))
elif item is not None:
result.append(item)
return result
def prepare_indexed_inputs(
main_data: Any,
*secondary_inputs: Any,
is_vector: list[bool] | None = None,
input_names: list[str] | None = None,
) -> tuple[bool, list[Any], list[list[Any]]]:
"""Prepares and validates main and secondary inputs according to the generalized list architecture:
- Main input can be a list or a single item.
- Secondary inputs that are lists must have either the same length as the main input, or be a single value (length 1 or non-list).
- If a secondary input is a single value, it is broadcast across all elements of the main input.
- If list lengths mismatch and length != 1, raises ValueError.
Returns:
(is_main_list, main_items, broadcasted_secondary_lists)
"""
main_is_list = is_input_a_list(main_data, is_vector=False)
if main_is_list:
main_items = list(main_data)
n = len(main_items)
else:
main_items = [main_data]
n = 1
vec_flags = is_vector or [False] * len(secondary_inputs)
names = input_names or [f"input_{i+1}" for i in range(len(secondary_inputs))]
broadcasted_secondaries: list[list[Any]] = []
for i, sec_val in enumerate(secondary_inputs):
is_vec = vec_flags[i] if i < len(vec_flags) else False
name = names[i] if i < len(names) else f"input_{i+1}"
sec_is_list = is_input_a_list(sec_val, is_vector=is_vec)
if not sec_is_list:
if main_is_list:
broadcasted_secondaries.append([sec_val] * n)
else:
broadcasted_secondaries.append([sec_val])
else:
m = len(sec_val)
if not main_is_list:
if m == 1:
broadcasted_secondaries.append([sec_val[0]])
else:
raise ValueError(
f"Input list length mismatch: main input is a single item, but secondary input '{name}' has {m} items (must be 1 or a single value)."
)
else:
if m == n:
broadcasted_secondaries.append(list(sec_val))
elif m == 1:
broadcasted_secondaries.append([sec_val[0]] * n)
else:
raise ValueError(
f"Input list length mismatch: main input has {n} items, but secondary input '{name}' has {m} items (must be {n} or 1)."
)
return main_is_list, main_items, broadcasted_secondaries
def eval_element_wise_multivalue(
op_func: Callable[..., Any],
main_data: Any,
*secondary_inputs: Any,
is_vector: list[bool] | None = None,
input_names: list[str] | None = None,
) -> Any:
"""Evaluates op_func element-wise with support for main and secondary list inputs:
- If main_data is a list, returns a list of results.
- If main_data is a single item, returns a single result.
- Secondary inputs that are lists are matched index per index.
- Secondary inputs that are single values (or length-1 lists) are broadcast across all main elements.
- Mismatched secondary list lengths raise ValueError.
"""
is_main_list, main_items, sec_lists = prepare_indexed_inputs(
main_data,
*secondary_inputs,
is_vector=is_vector,
input_names=input_names,
)
if is_main_list:
return [
op_func(item, *(sec[i] for sec in sec_lists))
for i, item in enumerate(main_items)
]
return op_func(main_items[0], *(sec[0] for sec in sec_lists))
def eval_element_wise(op_func: Callable[[Any], Any], data: Any) -> Any:
"""If data is a list, applies op_func to each element and returns a list.
Otherwise applies op_func to data directly."""
if is_list_data(data):
return [op_func(item) for item in data]
return op_func(data)
def set_socket_shape(sock: Any, shape: str) -> None:
"""Safely updates socket display_shape property ('LIST', 'CIRCLE', etc.)."""
if sock is not None:
try:
sock.display_shape = shape
except (AttributeError, TypeError):
pass
def _follow_reroute_upstream(sock: Any, visited: set | None = None) -> Any:
"""Follows reroute nodes backwards to find the original source output socket."""
if visited is None:
visited = set()
current = sock
while (
current is not None
and getattr(current, "node", None) is not None
and getattr(current.node, "bl_idname", "") == "NodeReroute"
and getattr(current.node, "inputs", None)
and current.node.inputs[0].is_linked
and current.node.inputs[0].links
):
ptr = current.as_pointer() if hasattr(current, "as_pointer") else id(current)
if ptr in visited:
break
visited.add(ptr)
current = current.node.inputs[0].links[0].from_socket
return current
def _follow_reroute_downstream(reroute_node: Any, visited: set | None = None) -> list[Any]:
"""Follows reroute nodes forwards to find all downstream target input sockets."""
if visited is None:
visited = set()
node_ptr = reroute_node.as_pointer() if hasattr(reroute_node, "as_pointer") else id(reroute_node)
if node_ptr in visited:
return []
visited.add(node_ptr)
targets: list[Any] = []
for outp in getattr(reroute_node, "outputs", []):
for link in getattr(outp, "links", []):
to_sock = link.to_socket
to_node = getattr(to_sock, "node", None)
if to_node and getattr(to_node, "bl_idname", "") == "NodeReroute":
targets.extend(_follow_reroute_downstream(to_node, visited))
else:
targets.append(to_sock)
return targets
def socket_accepts_list(sock: Any, visited: set | None = None) -> bool:
"""Determines whether a given input socket can accept a list input."""
if sock is None:
return False
if visited is None:
visited = set()
sock_ptr = sock.as_pointer() if hasattr(sock, "as_pointer") else id(sock)
if sock_ptr in visited:
return False
visited.add(sock_ptr)
node = getattr(sock, "node", None)
if node is None:
return False
bl_idname = getattr(node, "bl_idname", "")
# Strictly forbidden nodes
if bl_idname in ("SLICENODE_StartPrint", "SLICENODE_PrinterList"):
return False
# Socket type check
sock_type = getattr(sock, "bl_idname", "")
if sock_type not in LIST_CAPABLE_SOCKET_TYPES:
return False
# Reroute node: check downstream targets
if bl_idname == "NodeReroute":
targets = _follow_reroute_downstream(node, visited)
if not targets:
return False
return all(socket_accepts_list(t, visited) for t in targets)
# NodeGroupOutput
if bl_idname == "NodeGroupOutput":
return True
# Nested SLICENODE_Group
if bl_idname == "SLICENODE_Group":
inner_tree = getattr(node, "slicer_tree", None) or getattr(node, "node_tree", None)
if inner_tree:
return group_input_can_accept_list(inner_tree, sock, visited)
return False
return True
def _get_interface_item(tree: Any, sock_or_idx_or_ident: Any, in_out: str) -> Any:
attr_name = "inputs" if in_out == 'INPUT' else "outputs"
if not hasattr(tree, "interface"):
sockets = getattr(tree, attr_name, [])
if isinstance(sock_or_idx_or_ident, int) and 0 <= sock_or_idx_or_ident < len(sockets):
return sockets[sock_or_idx_or_ident]
ident = getattr(sock_or_idx_or_ident, "identifier", str(sock_or_idx_or_ident))
for s in sockets:
if s.identifier == ident or s.name == ident:
return s
return None
items = [
item for item in getattr(tree.interface, "items_tree", [])
if getattr(item, "item_type", "") == 'SOCKET' and getattr(item, "in_out", "") == in_out
]
if isinstance(sock_or_idx_or_ident, int):
return items[sock_or_idx_or_ident] if 0 <= sock_or_idx_or_ident < len(items) else None
ident = getattr(sock_or_idx_or_ident, "identifier", None)
name = getattr(sock_or_idx_or_ident, "name", None)
if isinstance(sock_or_idx_or_ident, str):
ident = name = sock_or_idx_or_ident
for item in items:
if ident and getattr(item, "identifier", None) == ident:
return item
for item in items:
if name and getattr(item, "name", None) == name:
return item
return None
def _get_input_interface_item(tree: Any, sock_or_idx_or_ident: Any) -> Any:
return _get_interface_item(tree, sock_or_idx_or_ident, 'INPUT')
def _get_output_interface_item(tree: Any, sock_or_idx_or_ident: Any) -> Any:
return _get_interface_item(tree, sock_or_idx_or_ident, 'OUTPUT')
def _find_group_input_socket(tree: Any, item: Any) -> Any:
if not item or not hasattr(tree, "nodes"):
return None
ident = getattr(item, "identifier", "")
name = getattr(item, "name", "")
for n in tree.nodes:
if getattr(n, "bl_idname", "") == "NodeGroupInput":
for outp in getattr(n, "outputs", []):
if outp.identifier == ident or outp.name == name:
return outp
return None
def _find_group_output_socket(tree: Any, item: Any) -> Any:
if not item or not hasattr(tree, "nodes"):
return None
ident = getattr(item, "identifier", "")
name = getattr(item, "name", "")
for n in tree.nodes:
if getattr(n, "bl_idname", "") == "NodeGroupOutput":
for inp in getattr(n, "inputs", []):
if inp.identifier == ident or inp.name == name:
return inp
return None
def group_input_can_accept_list(tree: Any, sock_or_idx_or_ident: Any, visited: set | None = None) -> bool:
"""Determines whether a group tree input socket can accept a list input."""
if not tree:
return False
if visited is None:
visited = set()
item = _get_input_interface_item(tree, sock_or_idx_or_ident)
if not item:
return False
item_key = (tree.as_pointer() if hasattr(tree, "as_pointer") else id(tree), getattr(item, "identifier", id(item)))
if item_key in visited:
return False
visited.add(item_key)
stype = getattr(item, "socket_type", "") or getattr(item, "bl_idname", "")
if stype not in LIST_CAPABLE_SOCKET_TYPES:
return False
gin_sock = _find_group_input_socket(tree, item)
if not gin_sock or not getattr(gin_sock, "is_linked", False) or not gin_sock.links:
return False
targets: list[Any] = []
for link in gin_sock.links:
to_sock = link.to_socket
to_node = getattr(to_sock, "node", None)
if to_node and getattr(to_node, "bl_idname", "") == "NodeReroute":
targets.extend(_follow_reroute_downstream(to_node, visited))
else:
targets.append(to_sock)
if not targets:
return False
return all(socket_accepts_list(t, visited) for t in targets)
def is_internal_socket_a_list(group_node: Any, tree: Any, in_sock: Any, visited: set | None = None) -> bool:
"""Determines whether an internal input socket receives a list within a group context."""
if not in_sock or not getattr(in_sock, "is_linked", False) or not in_sock.links:
return False
from_sock = in_sock.links[0].from_socket
from_sock = _follow_reroute_upstream(from_sock, visited)
if not from_sock:
return False
from_node = getattr(from_sock, "node", None)
if not from_node:
return False
if getattr(from_node, "is_list_output", False):
return True
if getattr(from_node, "bl_idname", "") == "NodeGroupInput":
in_item = _get_input_interface_item(tree, from_sock)
if not in_item:
return False
if group_node is not None:
for grp_in in getattr(group_node, "inputs", []):
if grp_in.identifier == getattr(in_item, "identifier", None) or grp_in.name == getattr(in_item, "name", None):
return is_socket_linked_to_list(grp_in)
return False
else:
return group_input_can_accept_list(tree, in_item, visited)
if getattr(from_node, "bl_idname", "") == "SLICENODE_Group":
inner = getattr(from_node, "slicer_tree", None) or getattr(from_node, "node_tree", None)
if inner:
return is_group_output_a_list(from_node, inner, from_sock, visited)
return False
if getattr(from_node, "bl_idname", "") == "SLICENODE_JoinGcode":
for inp in getattr(from_node, "inputs", []):
if getattr(inp, "name", "") == "Gcode" or getattr(inp, "identifier", "").startswith("Gcode_ms_in_"):
if is_internal_socket_a_list(group_node, tree, inp, visited):
return True
return False
mapping = getattr(from_node, "list_socket_mapping", None)
if mapping:
sock_name = getattr(from_sock, "name", "")
for m_in_name, m_out_names in mapping:
targets = [m_out_names] if isinstance(m_out_names, str) else list(m_out_names)
if sock_name in targets:
mapped_inp = getattr(from_node, "inputs", None) and from_node.inputs.get(m_in_name)
if mapped_inp:
return is_internal_socket_a_list(group_node, tree, mapped_inp, visited)
return getattr(from_sock, "display_shape", "") == "LIST"
def is_group_output_a_list(
group_node: Any,
tree: Any,
sock_or_idx_or_ident: Any,
visited: set | None = None
) -> bool:
"""Determines whether a group tree output socket produces a list."""
if not tree:
return False
if visited is None:
visited = set()
item = _get_output_interface_item(tree, sock_or_idx_or_ident)
if not item:
return False
item_key = (
tree.as_pointer() if hasattr(tree, "as_pointer") else id(tree),
getattr(item, "identifier", id(item)),
id(group_node) if group_node else 0,
)
if item_key in visited:
return False
visited.add(item_key)
gout_sock = _find_group_output_socket(tree, item)
if not gout_sock or not getattr(gout_sock, "is_linked", False) or not gout_sock.links:
return False
src_sock = gout_sock.links[0].from_socket
src_sock = _follow_reroute_upstream(src_sock, visited)
if not src_sock:
return False
src_node = getattr(src_sock, "node", None)
if not src_node:
return False
src_bl_id = getattr(src_node, "bl_idname", "")
if getattr(src_node, "is_list_output", False):
return True
if src_bl_id == "SLICENODE_JoinGcode":
for inp in getattr(src_node, "inputs", []):
if getattr(inp, "name", "") == "Gcode" or getattr(inp, "identifier", "").startswith("Gcode_ms_in_"):
if is_internal_socket_a_list(group_node, tree, inp, visited):
return True
return False
if src_bl_id == "SLICENODE_Group":
inner_tree = getattr(src_node, "slicer_tree", None) or getattr(src_node, "node_tree", None)
if inner_tree:
return is_group_output_a_list(src_node, inner_tree, src_sock, visited)
return False
if src_bl_id == "NodeGroupInput":
in_item = _get_input_interface_item(tree, src_sock)
if not in_item:
return False
if group_node is not None:
for grp_in in getattr(group_node, "inputs", []):
if grp_in.identifier == getattr(in_item, "identifier", None) or grp_in.name == getattr(in_item, "name", None):
return is_socket_linked_to_list(grp_in)
return False
else:
return group_input_can_accept_list(tree, in_item, visited)
mapping = getattr(src_node, "list_socket_mapping", None)
if mapping:
sock_name = getattr(src_sock, "name", "")
for in_name, out_names in mapping:
targets = [out_names] if isinstance(out_names, str) else list(out_names)
if sock_name in targets:
inp = getattr(src_node, "inputs", None) and src_node.inputs.get(in_name)
if inp:
return is_internal_socket_a_list(group_node, tree, inp, visited)
return getattr(src_sock, "display_shape", "") == "LIST"
def sync_group_node_socket_shapes(group_node: Any) -> None:
"""Updates display_shape property on a SLICENODE_Group node's inputs and outputs."""
tree = getattr(group_node, "slicer_tree", None) or getattr(group_node, "node_tree", None)
if not tree:
return
for inp in getattr(group_node, "inputs", []):
can_list = group_input_can_accept_list(tree, inp)
set_socket_shape(inp, "LIST" if can_list else "CIRCLE")
for outp in getattr(group_node, "outputs", []):
is_list = is_group_output_a_list(group_node, tree, outp)
set_socket_shape(outp, "LIST" if is_list else "CIRCLE")
def sync_group_tree_internal_sockets(tree: Any) -> None:
"""Updates display_shape property on NodeGroupInput and NodeGroupOutput within a group tree."""
if not tree or not hasattr(tree, "nodes"):
return
for node in tree.nodes:
bl_id = getattr(node, "bl_idname", "")
if bl_id == "NodeGroupInput":
for outp in getattr(node, "outputs", []):
if outp.name == "" or outp.identifier == "__extend__":
continue
can_list = group_input_can_accept_list(tree, outp)
set_socket_shape(outp, "LIST" if can_list else "CIRCLE")
elif bl_id == "NodeGroupOutput":
for inp in getattr(node, "inputs", []):
if inp.name == "" or inp.identifier == "__extend__":
continue
is_list = is_group_output_a_list(None, tree, inp)
set_socket_shape(inp, "LIST" if is_list else "CIRCLE")
def is_output_socket_a_list(sock: Any, visited: set | None = None) -> bool:
"""Determines whether an output socket produces a list."""
if sock is None:
return False
if visited is None:
visited = set()
sock_ptr = sock.as_pointer() if hasattr(sock, "as_pointer") else id(sock)
if sock_ptr in visited:
return False
visited.add(sock_ptr)
if getattr(sock, "display_shape", "") == "LIST":
return True
node = getattr(sock, "node", None)
if node is None:
return False
bl_idname = getattr(node, "bl_idname", "")
if getattr(node, "is_list_output", False):
return True
if bl_idname in ("SLICENODE_FloatMath", "SLICENODE_IntMath"):
return any(is_socket_linked_to_list(inp) for inp in getattr(node, "inputs", []))
if bl_idname == "NodeReroute":
inputs = getattr(node, "inputs", None)
if inputs and inputs[0].is_linked and inputs[0].links:
return is_output_socket_a_list(inputs[0].links[0].from_socket, visited)
return False
if bl_idname == "SLICENODE_JoinGcode":
inputs = getattr(node, "inputs", [])
return any(
is_socket_linked_to_list(inp)
for inp in inputs
if getattr(inp, "name", "") == "Gcode" or getattr(inp, "identifier", "").startswith("Gcode_ms_in_")
)
if bl_idname == "SLICENODE_Group":
tree = getattr(node, "slicer_tree", None) or getattr(node, "node_tree", None)
if tree:
return is_group_output_a_list(node, tree, sock, visited)
return False
if bl_idname == "NodeGroupInput":
tree = getattr(node, "id_data", None)
if tree:
return group_input_can_accept_list(tree, sock, visited)
return False
mapping = getattr(node, "list_socket_mapping", None)
if mapping:
sock_name = getattr(sock, "name", "")
for in_name, out_names in mapping:
targets = [out_names] if isinstance(out_names, str) else list(out_names)
if sock_name in targets:
inp = getattr(node, "inputs", None) and node.inputs.get(in_name)
if inp and is_socket_linked_to_list(inp):
return True
return False
def is_socket_linked_to_list(sock: Any) -> bool:
"""Determines whether an input socket is linked to an upstream list output."""
if not getattr(sock, "is_linked", False) or not getattr(sock, "links", None):
return False
from_socket = sock.links[0].from_socket
from_socket = _follow_reroute_upstream(from_socket)
return is_output_socket_a_list(from_socket)
def sync_node_list_sockets(node: Any, mapping: list[tuple[str, list[str] | str]]) -> None:
"""Updates display_shape on input and output sockets based on mapping.
Also updates display_shape for any secondary input sockets linked to a list.
mapping is a list of tuples: (input_socket_name, output_socket_names)
e.g. [('Paths', ['Paths']), ('Mesh', ['Selection', 'Inverted'])]
"""
mapped_inputs = set()
for in_name, out_names in mapping:
mapped_inputs.add(in_name)
inp = getattr(node, "inputs", None) and node.inputs.get(in_name)
if inp is None:
continue
is_list = is_socket_linked_to_list(inp)
set_socket_shape(inp, "LIST" if is_list else "CIRCLE")
if isinstance(out_names, str):
targets = [out_names]
else:
targets = list(out_names)
for out_name in targets:
out_sock = getattr(node, "outputs", None) and node.outputs.get(out_name)
if out_sock is not None:
set_socket_shape(out_sock, "LIST" if is_list else "CIRCLE")
for inp in getattr(node, "inputs", []):
if getattr(inp, "name", "") not in mapped_inputs:
is_list = is_socket_linked_to_list(inp)
set_socket_shape(inp, "LIST" if is_list else "CIRCLE")