velai/graph_undo.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from velai.dataflow.graph import DataGraph
from velai.nodes.base_node import BaseNode


@dataclass
class GraphUndoStack:
    """Stores snapshots of deleted graph fragments for Ctrl+Z restore."""

    max_entries: int = 50
    _stack: list[dict[str, Any]] = field(default_factory=list)

    def record_node_deletion(self, graph: DataGraph, node_ids: set[str]) -> None:
        if not node_ids:
            return

        nodes_data: list[dict[str, Any]] = []
        for node_id in node_ids:
            node = graph.nodes.get(node_id)
            if node is None:
                continue

            values: dict[str, Any] = {}
            if isinstance(node, BaseNode):
                try:
                    values = node.get_state()
                except Exception:
                    values = {}

            nodes_data.append(
                {
                    "id": node.node_id,
                    "kind": node.node_type.kind.value,
                    "x": node.x,
                    "y": node.y,
                    "values": values,
                    "width": node.width,
                    "height": node.height,
                }
            )

        if not nodes_data:
            return

        edges_data: list[dict[str, Any]] = []
        for conn in graph.connections:
            src_id = conn.start_node.node_id
            tgt_id = conn.end_node.node_id
            if src_id not in node_ids and tgt_id not in node_ids:
                continue
            edges_data.append(
                {
                    "source": src_id,
                    "sourceHandle": conn.start_port.name,
                    "target": tgt_id,
                    "targetHandle": conn.end_port.name,
                }
            )

        self._stack.append({"nodes": nodes_data, "edges": edges_data})
        if len(self._stack) > self.max_entries:
            self._stack.pop(0)

    def pop_deletion(self) -> dict[str, Any] | None:
        if not self._stack:
            return None
        return self._stack.pop()

    def clear(self) -> None:
        self._stack.clear()