velai/dataflow/codecs.py

from __future__ import annotations

from typing import Any

from .connection import Connection
from .nodes import NodeInstance
from .ports import PortSchema


def node_to_vueflow(n: NodeInstance, data_extra: dict[str, Any] | None = None) -> dict[str, Any]:
    """Convert a NodeInstance to a plain Vue Flow node dict.

    data_extra can be used by higher level code to attach UI specific state.
    If data_extra contains 'inputs' or 'outputs', they will replace the defaults.
    """
    data: dict[str, Any] = {
        "kind": n.node_type.kind.value,
        "title": n.node_type.display_name,
        "inputs": [_schema_json(p) for p in n.node_type.inputs],
        "outputs": [_schema_json(p) for p in n.node_type.outputs],
    }
    if data_extra:
        # app level UI code can attach arbitrary extra fields here
        # If inputs/outputs are provided, they replace the defaults (allows custom positioning)
        if "inputs" in data_extra:
            data["inputs"] = data_extra["inputs"]
        if "outputs" in data_extra:
            data["outputs"] = data_extra["outputs"]
        # Update other fields
        for key, value in data_extra.items():
            if key not in ("inputs", "outputs"):
                data[key] = value

    node_dict = {
        "id": n.node_id,
        "type": "base",
        "position": {"x": n.x, "y": n.y},
        "width": n.width,
        "height": n.height,
        "data": data,
    }

    return node_dict


def connection_to_vueflow_edge(c: Connection) -> dict[str, Any]:
    """Convert a Connection object to a Vue Flow edge dict."""
    source_handle = f"{c.start_node.node_id}:{c.start_port.name}"
    target_handle = f"{c.end_node.node_id}:{c.end_port.name}"

    edge_id = f"{c.start_node.node_id}:{c.start_port.name}->{c.end_node.node_id}:{c.end_port.name}"

    return {
        "id": edge_id,
        "source": c.start_node.node_id,
        "target": c.end_node.node_id,
        "sourceHandle": source_handle,
        "targetHandle": target_handle,
    }


def _schema_json(p: PortSchema) -> dict[str, Any]:
    return {
        "name": p.name,
        "type_id": p.dtype.id.value,
        "direction": p.direction.value,
        "multiplicity": p.multiplicity.value,
        "capacity": p.capacity,
        "color": p.dtype.color,
        "hidden": p.hidden,
        "tooltip": p.tooltip,
    }