|
| 1 | +"""Lay out a built node tree in columns that flow left to right. |
| 2 | +
|
| 3 | +Blender only measures nodes when it draws them, so builds in background mode |
| 4 | +or before the editor opens have no sizes to work with. Heights are estimated |
| 5 | +from each node's visible sockets and settings instead. |
| 6 | +""" |
| 7 | + |
| 8 | +from __future__ import annotations |
| 9 | + |
| 10 | +from collections import defaultdict |
| 11 | +from statistics import mean |
| 12 | + |
| 13 | +import bpy |
| 14 | + |
| 15 | +HEADER_HEIGHT = 30.0 |
| 16 | +ROW_HEIGHT = 22.0 |
| 17 | +COLUMN_GAP = 80.0 |
| 18 | +NODE_GAP = 40.0 |
| 19 | + |
| 20 | +# Unlinked inputs that show their value as a column of fields: one row per component. |
| 21 | +_EXPANDED_INPUT_ROWS = {"VECTOR": 4, "ROTATION": 4, "MATRIX": 1} |
| 22 | + |
| 23 | + |
| 24 | +def _visible(sockets: bpy.types.bpy_prop_collection) -> list[bpy.types.NodeSocket]: |
| 25 | + return [socket for socket in sockets if socket.enabled and not socket.hide] |
| 26 | + |
| 27 | + |
| 28 | +def _input_rows(socket: bpy.types.NodeSocket) -> int: |
| 29 | + if socket.is_linked or socket.hide_value: |
| 30 | + return 1 |
| 31 | + return _EXPANDED_INPUT_ROWS.get(socket.type, 1) |
| 32 | + |
| 33 | + |
| 34 | +def _base_properties(node: bpy.types.Node) -> set[str]: |
| 35 | + return { |
| 36 | + prop.identifier |
| 37 | + for base in type(node).__mro__[1:] |
| 38 | + if hasattr(base, "bl_rna") |
| 39 | + for prop in base.bl_rna.properties |
| 40 | + } |
| 41 | + |
| 42 | + |
| 43 | +def _setting_rows(node: bpy.types.Node) -> int: |
| 44 | + """Roughly how many rows the node's own settings take above its inputs.""" |
| 45 | + inherited = _base_properties(node) |
| 46 | + return sum( |
| 47 | + 1 |
| 48 | + for prop in node.bl_rna.properties |
| 49 | + if prop.identifier not in inherited |
| 50 | + and not prop.is_readonly |
| 51 | + and prop.type in {"ENUM", "BOOLEAN", "POINTER", "INT", "FLOAT", "STRING"} |
| 52 | + ) |
| 53 | + |
| 54 | + |
| 55 | +class _Measure: |
| 56 | + """Estimated sizes and socket offsets (downward from the node's top) of one node.""" |
| 57 | + |
| 58 | + def __init__(self, node: bpy.types.Node) -> None: |
| 59 | + outputs = _visible(node.outputs) |
| 60 | + inputs = _visible(node.inputs) |
| 61 | + self.width = node.width |
| 62 | + if node.hide: |
| 63 | + self.height = HEADER_HEIGHT |
| 64 | + self.output_offsets = {socket: HEADER_HEIGHT / 2 for socket in outputs} |
| 65 | + self.input_offsets = {socket: HEADER_HEIGHT / 2 for socket in inputs} |
| 66 | + return |
| 67 | + y = HEADER_HEIGHT |
| 68 | + self.output_offsets = {} |
| 69 | + for socket in outputs: |
| 70 | + self.output_offsets[socket] = y + ROW_HEIGHT / 2 |
| 71 | + y += ROW_HEIGHT |
| 72 | + y += _setting_rows(node) * ROW_HEIGHT |
| 73 | + self.input_offsets = {} |
| 74 | + for socket in inputs: |
| 75 | + self.input_offsets[socket] = y + ROW_HEIGHT / 2 |
| 76 | + y += _input_rows(socket) * ROW_HEIGHT |
| 77 | + self.height = y + ROW_HEIGHT / 2 |
| 78 | + |
| 79 | + |
| 80 | +def _consumers(node_tree: bpy.types.NodeTree) -> dict[bpy.types.Node, list[tuple[bpy.types.Node, bpy.types.NodeLink]]]: |
| 81 | + """Each node's downstream nodes, with the link to each (``None`` for a zone pairing).""" |
| 82 | + consumers = defaultdict(list) |
| 83 | + for link in node_tree.links: |
| 84 | + if link.from_node != link.to_node: |
| 85 | + consumers[link.from_node].append((link.to_node, link)) |
| 86 | + for node in node_tree.nodes: |
| 87 | + # Keep a zone's input to the left of its output even when nothing links them. |
| 88 | + paired = getattr(node, "paired_output", None) |
| 89 | + if paired is not None: |
| 90 | + consumers[node].append((paired, None)) |
| 91 | + return consumers |
| 92 | + |
| 93 | + |
| 94 | +def _columns( |
| 95 | + nodes: list[bpy.types.Node], consumers: dict[bpy.types.Node, list] |
| 96 | +) -> list[list[bpy.types.Node]]: |
| 97 | + """Group nodes by their longest path to a node nothing reads, rightmost column first.""" |
| 98 | + column: dict[bpy.types.Node, int] = {} |
| 99 | + visiting: set[bpy.types.Node] = set() |
| 100 | + |
| 101 | + def place(start: bpy.types.Node) -> None: |
| 102 | + stack = [(start, False)] |
| 103 | + while stack: |
| 104 | + node, expanded = stack.pop() |
| 105 | + if node in column or (not expanded and node in visiting): |
| 106 | + continue |
| 107 | + downstream = [ |
| 108 | + target for target, _ in consumers[node] if target not in visiting |
| 109 | + ] |
| 110 | + if not expanded: |
| 111 | + visiting.add(node) |
| 112 | + stack.append((node, True)) |
| 113 | + stack.extend((target, False) for target in downstream if target not in column) |
| 114 | + continue |
| 115 | + visiting.discard(node) |
| 116 | + column[node] = max((column[target] + 1 for target in downstream if target in column), default=0) |
| 117 | + |
| 118 | + for node in nodes: |
| 119 | + place(node) |
| 120 | + columns: list[list[bpy.types.Node]] = [[] for _ in range(max(column.values(), default=-1) + 1)] |
| 121 | + for node in nodes: |
| 122 | + columns[column[node]].append(node) |
| 123 | + return columns |
| 124 | + |
| 125 | + |
| 126 | +def arrange(node_tree: bpy.types.NodeTree) -> None: |
| 127 | + """Position the nodes of ``node_tree`` so links flow left to right without overlapping.""" |
| 128 | + nodes = [node for node in node_tree.nodes if node.bl_idname != "NodeFrame"] |
| 129 | + if not nodes: |
| 130 | + return |
| 131 | + consumers = _consumers(node_tree) |
| 132 | + measures = {node: _Measure(node) for node in nodes} |
| 133 | + top: dict[bpy.types.Node, float] = {} |
| 134 | + |
| 135 | + x = 0.0 |
| 136 | + for column in _columns(nodes, consumers): |
| 137 | + x -= max(measures[node].width for node in column) |
| 138 | + |
| 139 | + # Line each node's first linked output up with the inputs it feeds. |
| 140 | + desired = {} |
| 141 | + for node in column: |
| 142 | + targets = [ |
| 143 | + top[target] - measures[target].input_offsets.get(link.to_socket, 0.0) |
| 144 | + + measures[node].output_offsets.get(link.from_socket, 0.0) |
| 145 | + for target, link in consumers[node] |
| 146 | + if link is not None and target in top |
| 147 | + ] |
| 148 | + desired[node] = mean(targets) if targets else None |
| 149 | + |
| 150 | + # Unlinked nodes keep build order, after the linked ones. |
| 151 | + order = sorted( |
| 152 | + range(len(column)), |
| 153 | + key=lambda index: (desired[column[index]] is None, -(desired[column[index]] or 0.0), index), |
| 154 | + ) |
| 155 | + y = None |
| 156 | + offsets = [] |
| 157 | + for index in order: |
| 158 | + node = column[index] |
| 159 | + wanted = desired[node] if desired[node] is not None else (y if y is not None else 0.0) |
| 160 | + placed = wanted if y is None else min(wanted, y) |
| 161 | + top[node] = placed |
| 162 | + offsets.append(wanted - placed) |
| 163 | + y = placed - measures[node].height - NODE_GAP |
| 164 | + # Pushing nodes down to avoid overlaps moves them off their links; split the difference. |
| 165 | + shift = mean(offsets) |
| 166 | + for node in column: |
| 167 | + top[node] += shift |
| 168 | + node.location = (x, top[node]) |
| 169 | + x -= COLUMN_GAP |
| 170 | + |
| 171 | + # Start the tree at the origin. |
| 172 | + left = min(node.location.x for node in nodes) |
| 173 | + for node in nodes: |
| 174 | + node.location.x -= left |
0 commit comments