Copy/paste utilities for Nuke that preserve and restore node input connections.

copypaste_with_input_tracking.py

"""
https://gist.github.com/artandmath/22d9bf83128c7a73cca3136506c4d559

Copy/paste utilities for Nuke that preserve and restore node input connections.

When copying nodes, input connections are embedded as metadata in a hidden knob.
When pasting, that metadata is read and the original connections are re-established
(where possible), then the tracking knobs are removed.

Menu integration is available in two styles:
- Style A: Adds ``Edit > Copy For Input Restore`` (Ctrl+Alt+C) and overrides
  the standard ``Edit > Paste`` (Ctrl+V) with the restoring variant.
  If Edit > Copy is performed, then ``Edit > Paste`` functions as normal.
- Style B: Overrides ``Ctrl+C`` with the tracking copy, overrides
  ``Ctrl+V`` with :func:`paste_without_input_restore` (strips knobs, no reconnection),
  and adds ``Edit > Paste And Restore Inputs`` on ``Alt+Shift+V``.

Style A -- SAFER
The default and is recommended for most users because there is intent
by the user to restore the input connections.

Style B -- RISKIER
More natural "Excel style" copy/paste menu, but requires facility-wide installation to
reduce risk of script artifacts. Unlike Style A, every copy embeds tracking knobs —
not just intentional ones. If you install this in your personal .nuke while others use
standard copy/paste, sharing node snippets risks leaking tracking knobs into their
scripts. Think before you paste!

OPTIONAL POWER USER FEATURE: Enable at your own confusion.

When ``PASTE_TO_SELECTED_MENU`` is ``True``, both styles also register:
- ``Edit > Paste To Selected Nodes`` (Ctrl+Shift+V) — pastes the clipboard once per
  originally-selected node, connecting the pasted nodes to each target in turn.
- ``Edit > Paste To Selected Nodes And Restore Inputs`` (Ctrl+Alt+Shift+V) — same as
  above but also restores original input connections after each individual paste.
"""

MENU_STYLE = "A"
PASTE_TO_SELECTED_MENU = False

import re
import tempfile

import nuke
from PySide2 import QtWidgets


def copy_for_input_restore():
    """Copy selected nodes to the clipboard with input-connection metadata embedded.

    For each selected node, the indices and names of its connected inputs are
    serialised into a hidden ``addUserKnob`` entry injected into the ``.nk`` text
    before it is placed on the system clipboard.  A subsequent call to
    :func:`paste_and_restore_inputs` can then read that metadata and reconnect the
    nodes automatically.

    Does nothing when no nodes are selected.
    """
    sel = nuke.selectedNodes()
    if not sel:
        return

    input_map = {}
    for node in sel:
        inputs = []
        for i in range(node.inputs()):
            inp = node.input(i)
            if inp:
                inputs.append(f"{i}:{inp.name()}")
        input_map[node.name()] = "\\n".join(inputs)

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".nk")
    path = tmp.name
    tmp.close()

    nuke.nodeCopy(path)

    with open(path, "r") as f:
        nk = f.read()

    lines = nk.split("\n")
    node_name = None
    new_lines = []

    for line in lines:
        name_match = re.match(r"\s*name\s+(\S+)", line)
        if name_match:
            node_name = name_match.group(1)

        if line.strip() == "}" and node_name in input_map:
            inputs_text = input_map[node_name]
            new_lines.append(" addUserKnob {20 original_inputs_tracking_tab}")
            new_lines.append(
                f' addUserKnob {{26 original_inputs_tracking_knob -STARTLINE T "{inputs_text}"}}'
            )
            node_name = None

        new_lines.append(line)

    clipboard = QtWidgets.QApplication.clipboard()
    clipboard.setText("\n".join(new_lines))


def _strip_tracking_knobs(node):
    """Remove input-tracking knobs from *node* if present.

    Removes ``original_inputs_tracking_tab`` first (so the tab is gone before
    the child knob), then ``original_inputs_tracking_knob``.  Safe to call on
    nodes that carry neither knob.
    """
    if "original_inputs_tracking_knob" in node.knobs():
        node.removeKnob(node["original_inputs_tracking_knob"])
    if "original_inputs_tracking_tab" in node.knobs():
        node.removeKnob(node["original_inputs_tracking_tab"])


def _restore_inputs(nodes):
    """Restore input connections from tracking metadata for a list of *nodes*.

    For each node that carries an ``original_inputs_tracking_knob`` knob, the
    encoded connection data is parsed and each input is reconnected to the named
    source node (if that node exists in the current script and the input slot is
    not already occupied).  Tracking knobs are removed after processing.
    """
    for node in nodes:
        if "original_inputs_tracking_knob" not in node.knobs():
            _strip_tracking_knobs(node)
            continue

        data = node["original_inputs_tracking_knob"].value().strip()
        if not data:
            _strip_tracking_knobs(node)
            continue

        for line in data.split("\n"):
            try:
                idx_str, input_name = line.split(":", 1)
                idx = int(idx_str)
            except ValueError:
                continue

            if node.input(idx):
                continue

            target = nuke.toNode(input_name)
            if target:
                node.setInput(idx, target)

        _strip_tracking_knobs(node)


def paste_and_restore_inputs():
    """Paste nodes from the clipboard and restore their original input connections.

    Deselects all existing nodes, performs a standard paste, then delegates to
    :func:`_restore_inputs` to reconnect inputs and strip tracking knobs from
    every pasted node.

    Does nothing when the clipboard yields no new nodes.
    """
    for n in nuke.allNodes():
        n["selected"].setValue(False)

    nuke.nodePaste("%clipboard%")

    pasted = nuke.selectedNodes()
    if not pasted:
        return

    _restore_inputs(pasted)


def paste_without_input_restore():
    """Paste nodes from the clipboard and strip input-tracking knobs without reconnecting.

    Behaves identically to :func:`paste_and_restore_inputs` except that no
    attempt is made to restore input connections.  Use this when you deliberately
    want the pasted nodes to be disconnected but still need the temporary tracking
    knobs cleaned up.

    Does nothing when the clipboard yields no new nodes.
    """
    for n in nuke.allNodes():
        n["selected"].setValue(False)

    nuke.nodePaste("%clipboard%")

    pasted = nuke.selectedNodes()
    if not pasted:
        return

    for node in pasted:
        _strip_tracking_knobs(node)


def paste_to_selected():
    """Paste clipboard nodes once per originally-selected node, then re-select all pasted nodes.

    Captures the current selection, deselects everything, then iterates over
    each originally-selected node.  For each one, that node is selected in
    isolation and a paste is performed so that Nuke connects the incoming nodes
    to it.  All pasted nodes are collected across every iteration.  Once the
    loop is complete, only the pasted nodes are selected.

    Tracking knobs left by :func:`copy_for_input_restore` are removed from
    every pasted node.

    Does nothing when no nodes are selected.
    """
    original_selection = nuke.selectedNodes()
    if not original_selection:
        return

    for n in nuke.allNodes():
        n["selected"].setValue(False)

    all_pasted = []

    for target in original_selection:
        target["selected"].setValue(True)

        nuke.nodePaste("%clipboard%")

        pasted = nuke.selectedNodes()
        all_pasted.extend(pasted)

        for n in nuke.allNodes():
            n["selected"].setValue(False)

    for node in all_pasted:
        _strip_tracking_knobs(node)
        node["selected"].setValue(True)


def paste_to_selected_and_restore_inputs():
    """Paste clipboard nodes once per originally-selected node and restore input connections.

    Identical to :func:`paste_to_selected` except that after each individual
    paste the input-connection metadata embedded by :func:`copy_for_input_restore`
    is read and the connections are re-established before moving on to the next
    target node.

    Does nothing when no nodes are selected.
    """
    original_selection = nuke.selectedNodes()
    if not original_selection:
        return

    for n in nuke.allNodes():
        n["selected"].setValue(False)

    all_pasted = []

    for target in original_selection:
        target["selected"].setValue(True)

        nuke.nodePaste("%clipboard%")

        pasted = nuke.selectedNodes()
        _restore_inputs(pasted)
        all_pasted.extend(pasted)

        for n in nuke.allNodes():
            n["selected"].setValue(False)

    for node in all_pasted:
        node["selected"].setValue(True)


# ---------------------------------------------------------------------------
# Menu integration
# ---------------------------------------------------------------------------

def _register_menu_style_a():
    """Register Style A menu commands.

    Always adds:
    - ``Edit > Copy For Input Restore`` (Ctrl+Alt+C) — copy with tracking metadata.
    - ``Edit > Paste`` (Ctrl+V) — overrides the default paste with input restore.

    When ``PASTE_TO_SELECTED_MENU`` is ``True``, also adds:
    - ``Edit > Paste To Selected Nodes`` (Ctrl+Shift+V)
    - ``Edit > Paste To Selected Nodes And Restore Inputs`` (Ctrl+Alt+Shift+V)
    """
    menu = nuke.menu("Nuke")
    edit_menu = menu.findItem("Edit")

    copy_index = None
    paste_index = None
    for i, item in enumerate(edit_menu.items()):
        name = item.name()
        if name == "Copy":
            copy_index = i
        elif name == "Paste":
            paste_index = i

    if copy_index is not None:
        menu.addCommand(
            "Edit/Copy For Input Restore",
            "copy_for_input_restore()",
            "Ctrl+Alt+C",
            index=copy_index + 1,
        )

    if paste_index is not None:
        menu.addCommand(
            "Edit/Paste",
            "paste_and_restore_inputs()",
            "Ctrl+V",
        )

    if PASTE_TO_SELECTED_MENU:
        menu.addCommand(
            "Edit/Paste To Selected Nodes",
            "paste_to_selected()",
            "Ctrl+Shift+V",
            index=paste_index + 2,
        )
        menu.addCommand(
            "Edit/Paste To Selected Nodes And Restore Inputs",
            "paste_to_selected_and_restore_inputs()",
            "Ctrl+Alt+Shift+V",
            index=paste_index + 3,
        )


def _register_menu_style_b():
    """Register Style B menu commands.

    Always adds:
    - ``Edit > Copy`` (Ctrl+C) — overrides the default copy with tracking metadata.
    - ``Edit > Paste`` (Ctrl+V) — overrides the default paste; strips tracking knobs
      but does not restore connections.
    - ``Edit > Paste And Restore Inputs`` (Alt+Shift+V) — paste with full reconnection.

    When ``PASTE_TO_SELECTED_MENU`` is ``True``, also adds:
    - ``Edit > Paste To Selected Nodes`` (Ctrl+Shift+V)
    - ``Edit > Paste To Selected Nodes And Restore Inputs`` (Ctrl+Alt+Shift+V)
    """
    menu = nuke.menu("Nuke")
    edit_menu = menu.findItem("Edit")

    copy_index = None
    paste_index = None
    for i, item in enumerate(edit_menu.items()):
        name = item.name()
        if name == "Copy":
            copy_index = i
        elif name == "Paste":
            paste_index = i

    if copy_index is not None:
        menu.addCommand(
            "Edit/Copy",
            "copy_for_input_restore()",
            "Ctrl+C",
        )

    if PASTE_TO_SELECTED_MENU:
        menu.addCommand(
            "Edit/Paste To Selected Nodes",
            "paste_to_selected()",
            "Ctrl+Shift+V",
            index=paste_index + 1,
        )
        menu.addCommand(
            "Edit/Paste To Selected Nodes And Restore Inputs",
            "paste_to_selected_and_restore_inputs()",
            "Ctrl+Alt+Shift+V",
            index=paste_index + 2,
        )

    if paste_index is not None:
        menu.addCommand(
            "Edit/Paste",
            "paste_without_input_restore()",
            "Ctrl+V",
        )
        menu.addCommand(
            "Edit/Paste And Restore Inputs",
            "paste_and_restore_inputs()",
            "Alt+Shift+V",
            index=paste_index + (3 if PASTE_TO_SELECTED_MENU else 1),
        )

if MENU_STYLE == "A":
    _register_menu_style_a()
elif MENU_STYLE == "B":
    _register_menu_style_b()
else:
    raise ValueError(f"Unknown MENU_STYLE {MENU_STYLE!r}; expected 'A' or 'B'.")