python_code_node.py

"""Python-based custom node utilities for ComfyUI."""

from __future__ import annotations

import difflib
import io
import traceback
from contextlib import redirect_stdout
from pathlib import Path
from typing import Any, Dict, List, Tuple

try:  # pragma: no cover - ComfyUI runtime provides these modules
    from aiohttp import web
    from server import PromptServer
except Exception:  # pragma: no cover - tests may run without ComfyUI
    web = None
    PromptServer = None


def _stringify_result_element(value: Any, delimiter: str) -> str:
    if isinstance(value, (list, tuple)):
        parts = [_stringify_result_element(item, delimiter) for item in value]
        joiner = delimiter if delimiter is not None else ""
        return joiner.join(parts)
    return str(value)


class PythonCodeNode:
    """Execute Python code with helpers for working with ComfyUI strings."""

    CATEGORY = "utils/code"
    FUNCTION = "run"
    RETURN_TYPES = ("STRING", "STRING", "LIST", "STRING", "STRING", "BOOLEAN")
    RETURN_NAMES = ("result", "result_lines", "result_lines_list", "stdout", "stderr", "ok")
    OUTPUT_IS_LIST = (False, False, True, False, False, False)
    INPUT_IS_LIST = False
    MAX_INPUT_SLOTS = 20
    DEFAULT_INPUT_SLOTS = 1
    EXTENSION_ROOT = Path(__file__).resolve().parent

    @classmethod
    def INPUT_TYPES(cls):
        def multiline_str(name: str):
            index = int(name.replace("input", ""))
            alias = "inputs[0]" if index == 1 else f"inputs[{index - 1}]"
            alias_hint = " (also input_text/lines)" if index == 1 else ""
            placeholder = f"{name} ↔ {alias}{alias_hint}; string/list depends on split_lines"
            return {
                "multiline": True,
                "default": "",
                "placeholder": placeholder,
            }

        optional_inputs: Dict[str, Tuple[str, Dict[str, Any]]] = {}
        for slot in range(2, cls.MAX_INPUT_SLOTS + 1):
            optional_inputs[f"input{slot}"] = ("STRING", multiline_str(f"input{slot}"))
        optional_inputs["load_from_file"] = ("BOOLEAN", {"default": False})
        optional_inputs["script_filename"] = (
            "STRING",
            {
                "default": "",
                "multiline": False,
                "placeholder": "Relative to this extension (e.g., scripts/example.py)",
            },
        )
        optional_inputs["split_lines"] = ("BOOLEAN", {"default": True})
        optional_inputs["strip_empty"] = ("BOOLEAN", {"default": True})
        optional_inputs["delimiter"] = (
            "STRING",
            {
                "default": ", ",
                "multiline": False,
                "placeholder": "Optional custom delimiter (comma, pipe, etc.)",
            },
        )
        optional_inputs["output_inner_delimiter"] = (
            "STRING",
            {
                "default": ", ",
                "multiline": False,
                "placeholder": "Join nested lists in auto result_lines (default matches delimiter).",
            },
        )
        optional_inputs["input_slots"] = (
            "INT",
            {
                "default": cls.DEFAULT_INPUT_SLOTS,
                "min": 1,
                "max": cls.MAX_INPUT_SLOTS,
                "step": 1,
                "display": "number",
            },
        )

        return {
            "required": {
                "script": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": "result = input1",
                        "placeholder": "Set result/result_lines; access raw data via inputs[*]. Additional inputs appear automatically as you type.",
                    },
                ),
                "input1": ("STRING", multiline_str("input1")),
            },
            "optional": optional_inputs,
        }

    def run(
        self,
        script: str,
        input1: str = "",
        input2: str = "",
        input3: str = "",
        input4: str = "",
        input5: str = "",
        input6: str = "",
        input7: str = "",
        input8: str = "",
        input9: str = "",
        input10: str = "",
        input11: str = "",
        input12: str = "",
        input13: str = "",
        input14: str = "",
        input15: str = "",
        input16: str = "",
        input17: str = "",
        input18: str = "",
        input19: str = "",
        input20: str = "",
        load_from_file: bool = False,
        script_filename: str = "",
        input_slots: int = DEFAULT_INPUT_SLOTS,
        split_lines: bool = True,
        strip_empty: bool = True,
        delimiter: str = ", ",
        output_inner_delimiter: str = ", ",
    ) -> Tuple[str, str, List[str], str, str, bool]:
        """Execute *script* and expose helpers for returning data to ComfyUI."""

        stdout_buffer = io.StringIO()
        stderr = ""
        ok = True
        result_text = ""
        result_lines_list: List[str] = []
        result_value: Any = None
        raw_inputs = [
            input1,
            input2,
            input3,
            input4,
            input5,
            input6,
            input7,
            input8,
            input9,
            input10,
            input11,
            input12,
            input13,
            input14,
            input15,
            input16,
            input17,
            input18,
            input19,
            input20,
        ]
        normalized_inputs: List[str] = [str(value or "") for value in raw_inputs[: self.MAX_INPUT_SLOTS]]
        input_line_sets: List[List[str]] = [text.splitlines() for text in normalized_inputs]
        delimiter_value = str(delimiter or "")
        use_delimiter = bool(delimiter_value)

        def split_text(value: str) -> List[str]:
            parts = [value]
            if use_delimiter and delimiter_value:
                new_parts: List[str] = []
                for chunk in parts:
                    new_parts.extend(chunk.split(delimiter_value))
                parts = new_parts
            if split_lines:
                new_parts = []
                for chunk in parts:
                    new_parts.extend(chunk.splitlines())
                parts = new_parts
            return parts

        split_inputs: List[List[str]] = [split_text(text) for text in normalized_inputs]

        script_source = script
        script_path_display = ""
        if load_from_file:
            filename = (script_filename or "").strip()
            if not filename:
                return (
                    "",
                    [],
                    "",
                    "load_from_file is enabled but no script_filename was provided.",
                    False,
                )
            script_path = Path(filename)
            if not script_path.is_absolute():
                script_path = (self.EXTENSION_ROOT / script_path).resolve()
            else:
                script_path = script_path.resolve()
            script_path_display = str(script_path)
            try:
                script_source = script_path.read_text(encoding="utf-8")
            except Exception as exc:  # pragma: no cover - relies on filesystem state
                return (
                    "",
                    [],
                    "",
                    f"Failed to load script from '{filename}': {exc}",
                    False,
                )

        try:
            requested_slots = int(input_slots)
        except (TypeError, ValueError):
            requested_slots = self.DEFAULT_INPUT_SLOTS
        active_inputs = max(1, min(self.MAX_INPUT_SLOTS, requested_slots))
        primary_input = normalized_inputs[0]
        primary_lines = input_line_sets[0]
        split_mode = split_lines or use_delimiter
        inputs_payload = split_inputs[:active_inputs] if split_mode else normalized_inputs[:active_inputs]
        local_ns: Dict[str, Any] = {
            "input_text": primary_input,
            "lines": primary_lines,
            "result": "",
            "result_lines": [],
            "inputs": inputs_payload,
            "inputs_text": normalized_inputs[:active_inputs],
            "inputs_lines": input_line_sets[:active_inputs],
            "active_inputs": active_inputs,
            "input_slots": active_inputs,
            "script_path": script_path_display,
            "delimiter": delimiter_value,
            "output_inner_delimiter": output_inner_delimiter,
        }

        for index, text_value in enumerate(normalized_inputs, start=1):
            lines = input_line_sets[index - 1]
            local_ns[f"input{index}_text"] = text_value
            local_ns[f"input{index}_lines"] = lines
            local_ns[f"input{index}"] = split_inputs[index - 1] if split_mode else text_value

        try:
            local_ns.setdefault("__builtins__", __builtins__)
            with redirect_stdout(stdout_buffer):
                exec(script_source, local_ns, local_ns)
            result_value = local_ns.get("result", None)
            if result_value is None and "result_text" in local_ns:
                result_value = local_ns.get("result_text")
            result_text = "" if result_value is None else str(result_value)
            result_lines_value = local_ns.get("result_lines", None)
            if result_lines_value is None:
                result_lines_list = []
            elif isinstance(result_lines_value, list):
                result_lines_list = result_lines_value
            else:
                result_lines_list = list(result_lines_value)
        except Exception:  # pragma: no cover - safety against runtime errors
            ok = False
            result_text = ""
            result_lines_list = []
            stderr = traceback.format_exc()
            result_value = None

        stdout = stdout_buffer.getvalue()

        if not result_lines_list and isinstance(result_value, (list, tuple)):
            result_lines_list = [
                _stringify_result_element(item, output_inner_delimiter) for item in result_value
            ]
        elif split_lines:
            auto_lines = result_text.splitlines()
            if strip_empty:
                auto_lines = [line for line in auto_lines if line.strip()]
            if not result_lines_list:
                result_lines_list = auto_lines

        result_lines_list = [
            _stringify_result_element(line, output_inner_delimiter) for line in result_lines_list
        ]
        if strip_empty:
            result_lines_list = [line for line in result_lines_list if line.strip()]
        result_lines_list_output = list(result_lines_list)
        result_lines_text = "\n".join(result_lines_list_output)

        return result_text, result_lines_text, result_lines_list_output, stdout, stderr, ok


def _resolve_script_destination(filename: str) -> Path:
    sanitized = (filename or "").strip()
    if not sanitized:
        raise ValueError("script filename is required")
    candidate = Path(sanitized)
    if not candidate.is_absolute():
        candidate = (PythonCodeNode.EXTENSION_ROOT / candidate).resolve()
    else:
        candidate = candidate.resolve()
    try:
        candidate.relative_to(PythonCodeNode.EXTENSION_ROOT)
    except ValueError as exc:
        raise ValueError("script_filename must remain inside the code-nodes directory") from exc
    return candidate


def _json_reply(ok: bool, message: str = "", status: int = 200, **extra):
    payload: Dict[str, Any] = {"ok": ok, "message": message}
    payload.update(extra)
    return web.json_response(payload, status=status)


def register_routes() -> None:
    if not (web and PromptServer and getattr(PromptServer, "instance", None)):
        return
    server = PromptServer.instance
    if getattr(server, "_code_nodes_routes", False):  # type: ignore[attr-defined]
        return

    @server.routes.get("/code-nodes/script")
    async def load_script(request):
        path_value = request.rel_url.query.get("path", "")
        try:
            destination = _resolve_script_destination(path_value)
        except ValueError as exc:
            return _json_reply(False, str(exc), status=400)

        if not destination.exists():
            return _json_reply(False, "File not found.", status=404)

        try:
            contents = destination.read_text(encoding="utf-8")
        except Exception as exc:  # pragma: no cover - filesystem
            return _json_reply(False, f"Failed to read file: {exc}", status=500)

        return _json_reply(True, "Loaded script.", path=str(destination), contents=contents)

    @server.routes.post("/code-nodes/script")
    async def save_script(request):
        try:
            data = await request.json()
        except Exception:
            return _json_reply(False, "Invalid JSON payload", status=400)

        path_value = data.get("path", "")
        contents = data.get("contents", "")
        force = bool(data.get("force"))
        try:
            destination = _resolve_script_destination(path_value)
        except ValueError as exc:
            return _json_reply(False, str(exc), status=400)

        if not isinstance(contents, str):
            return _json_reply(False, "contents must be a string", status=400)

        destination.parent.mkdir(parents=True, exist_ok=True)
        exists = destination.exists()
        text_before = ""
        if exists:
            text_before = destination.read_text(encoding="utf-8")
            if text_before == contents:
                return _json_reply(True, "File already up to date.", path=str(destination))
            if not force:
                diff_lines = difflib.unified_diff(
                    text_before.splitlines(),
                    contents.splitlines(),
                    fromfile=str(destination),
                    tofile="pending changes",
                    lineterm="",
                )
                diff_text = "\n".join(diff_lines).strip()
                if not diff_text:
                    diff_text = "(no textual diff available)"
                return _json_reply(
                    False,
                    "File exists and differs.",
                    requires_confirmation=True,
                    diff=diff_text,
                )

        destination.write_text(contents, encoding="utf-8")
        return _json_reply(True, "Script saved.", path=str(destination))

    server._code_nodes_routes = True  # type: ignore[attr-defined]


register_routes()