Fusion script to swap two selected tools, which swaps their position and connections

fusion_swap_tools.py

from typing import List, Iterable
import itertools
import contextlib


@contextlib.contextmanager
def comp_lock_and_undo_chunk(
    comp,
    undo_queue_name="Script CMD",
    keep_undo=True,
):
    """Lock comp and open an undo chunk during the context"""
    try:
        comp.Lock()
        comp.StartUndo(undo_queue_name)
        yield
    finally:
        comp.Unlock()
        comp.EndUndo(keep_undo)


def iter_inputs(tool):
    inputs = list(tool.GetInputList().values())
    for input in inputs:
        connected_output = input.GetConnectedOutput()
        if connected_output:
            yield connected_output, input

def iter_outputs(tool):
    outputs = list(tool.GetOutputList().values())
    for output in outputs:
        connected_inputs = list(output.GetConnectedInputs().values())
        for connected_input in connected_inputs:
            yield output, connected_input


def get_position(tool) -> List[float]:
    return list(
        tool.Composition.CurrentFrame.FlowView.GetPosTable(tool).values()
    )


def set_position(tool, pos: List[float]):
    tool.Composition.CurrentFrame.FlowView.SetPos(tool, *pos)


def swap_connections(a, b):
    a_inputs = list(iter_inputs(a))
    b_inputs = list(iter_inputs(b))
    a_outputs = list(iter_outputs(a))
    b_outputs = list(iter_outputs(b))

    # Disconnect all
    for output, input in itertools.chain(
            a_inputs, b_inputs, a_outputs, b_outputs):
        input.ConnectTo(None)

    # Reconnect from the other nodes
    for output, input in a_inputs:
        b_input = b[input.Name]
        if b_input:
            b_input.ConnectTo(output)

    for output, input in b_inputs:
        a_input = a[input.Name]
        if a_input:
            a_input.ConnectTo(output)

    for output, input in a_outputs:
        b_output = b[output.Name]
        if b_output:
            input.ConnectTo(b_output)

    for output, input in b_outputs:
        a_output = a[output.Name]
        if a_output:
            input.ConnectTo(a_output)


def swap_positions(a, b):
    a_pos = get_position(a)
    b_pos = get_position(b)
    set_position(a, b_pos)
    set_position(b, a_pos)


def swap_tools(a, b):
    """Swap tool a with tool b, swapping connections and positions."""
    swap_connections(a, b)
    swap_positions(a, b)


if __name__ == "__main__":
    # Selected tools
    tools = list(comp.GetToolList(True).values())
    if len(tools) != 2:
        raise ValueError("Please select two tools to swap")

    names = [tool.Name for tool in tools]
    print(f"Swapping tools: {tools[0].Name} <-> {tools[1].Name}")
    with comp_lock_and_undo_chunk(comp):
        swap_tools(tools[0], tools[1])