ComfyUI/custom_nodes/koolook/tools/preflight_release.py

#!/usr/bin/env python3
"""
Pre-flight release validation for ComfyUI-Koolook.

Runs the release-gate checks and exits non-zero on any failure:

  1. static       AST extraction of NODE_CLASS_MAPPINGS keys; verifies all
                  *.py files parse and that we can produce a definitive list
                  of registered node IDs without importing torch.
  2. dispatch     Stub-VAE round through the v2_3_3 wrapper; verifies all
                  five rank/VAE dispatch branches still wire correctly.
  3. workflows    Walks tests/workflows/*.json, extracts node IDs that look
                  like Koolook-pack IDs, verifies each one is currently
                  registered (i.e. wouldn't break a saved workflow).

Optional advisory check:

  manager-meta    Fetches ComfyUI-Manager's extension-node-map.json upstream
                  and diffs the entry for our repo against the AST-extracted
                  NODE_CLASS_MAPPINGS. This catches phantom-nodes /
                  missing-nodes drift, but does not block releases because
                  upstream Manager metadata is maintained independently.

Usage:
    python tools/preflight_release.py                # run release gates
    python tools/preflight_release.py --check static
    python tools/preflight_release.py --check workflows --verbose
    python tools/preflight_release.py --check manager-meta  # advisory only

Exit codes:
    0  all run checks passed
    1  one or more checks failed
    2  invalid CLI usage
"""
from __future__ import annotations

import argparse
import ast
import json
import re
import sys
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Iterable

REPO_ROOT = Path(__file__).resolve().parent.parent
REPO_URL = "https://github.com/malkuthro/ComfyUI-Koolook"
MANAGER_EXT_MAP_URL = (
    "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/"
    "main/extension-node-map.json"
)


class _HttpsOnlyRedirectHandler(urllib.request.HTTPRedirectHandler):
    """Refuses to follow a 30x redirect to a non-https target.

    `urllib.request.urlopen` follows redirects by default; a one-time
    https URL check at the call site doesn't protect against an
    intermediate hop redirecting to file:// or http:// (Bandit B310's
    real concern). This handler re-validates each redirect target.
    """

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        if not newurl.startswith("https://"):
            raise urllib.error.URLError(
                f"refusing redirect to non-https scheme: {newurl}"
            )
        return super().redirect_request(req, fp, code, msg, headers, newurl)


# ANSI escape codes (https://en.wikipedia.org/wiki/ANSI_escape_code).
# No-op when piped, redirected, or NO_COLOR is set.
_ANSI = {
    "red": "31",
    "green": "32",
    "yellow": "33",
    "cyan": "36",
    "dim": "2",
}


def _color(s: str, code: str) -> str:
    if not sys.stdout.isatty():
        return s
    return f"\033[{code}m{s}\033[0m"


def _green(s: str) -> str: return _color(s, _ANSI["green"])
def _red(s: str) -> str: return _color(s, _ANSI["red"])
def _yellow(s: str) -> str: return _color(s, _ANSI["yellow"])
def _cyan(s: str) -> str: return _color(s, _ANSI["cyan"])
def _dim(s: str) -> str: return _color(s, _ANSI["dim"])


# ───────────────────────────────────────────────────────────────────────
# Check 1 — static (AST extraction)
# ───────────────────────────────────────────────────────────────────────


# Files / folders to skip when walking for *.py.
_AST_SKIP_DIRS = {
    "__pycache__", ".venv", ".git", "upscaler_FIX",
    "nuke_CAM_exporter", "node_modules",
}

# Some fork implementation files retain upstream-style mappings so they can be
# inspected or imported in isolation. Runtime registration goes through the
# package entrypoint, where these IDs are namespaced or superseded.
_STATIC_NODE_ID_OMISSIONS = {
    Path("forks/whatdreamscost_koolook/versions/v1_3_2/ltx_director.py"): {
        "LTXDirector",
    },
}


def _iter_python_files() -> Iterable[Path]:
    for p in REPO_ROOT.rglob("*.py"):
        if any(part in _AST_SKIP_DIRS for part in p.parts):
            continue
        yield p


def _extract_literal_dict_keys(node: ast.AST) -> set[str]:
    """Pull literal string keys out of an ast.Dict node, ignoring computed keys."""
    if not isinstance(node, ast.Dict):
        return set()
    keys = set()
    for k in node.keys:
        if isinstance(k, ast.Constant) and isinstance(k.value, str):
            keys.add(k.value)
    return keys


def check_static(verbose: bool = False) -> tuple[bool, set[str]]:
    """
    Walk all .py files, AST-parse each, and collect every literal-dict
    `NODE_CLASS_MAPPINGS = {...}` source-side key. Returns (ok, ids).

    What this DOES NOT do (and the workflow check is the safety net):
      - Follow runtime-computed mappings (dict comprehensions, function
        calls like `_namespace_mappings(...)` in the v2_3_3 wrapper).
      - Apply `NAMESPACE_SUFFIX` / `SKIP_VERSION_SUFFIX` logic — the
        current code only walks `NODE_CLASS_MAPPINGS`, not the fork's
        suffix-mangling glue.

    For the Radiance Koolook v2_3_3 fork, this happens to work because
    every exposed VAE ID is in `SKIP_VERSION_SUFFIX`, so the source-dict
    keys match the user-facing IDs. The first time a non-skipped ID lands
    in a fork, the workflow-fixture check (4) will surface the mismatch
    by reporting the suffixed live ID as "missing" — at that point this
    function should grow SKIP_VERSION_SUFFIX extraction.
    """
    print(_cyan("static — AST extraction of NODE_CLASS_MAPPINGS"))
    parse_errors: list[str] = []
    all_node_ids: set[str] = set()
    files_seen = 0

    for py_path in _iter_python_files():
        files_seen += 1
        try:
            source = py_path.read_text(encoding="utf-8")
            tree = ast.parse(source, filename=str(py_path))
        except (SyntaxError, UnicodeDecodeError) as exc:
            parse_errors.append(f"{py_path.relative_to(REPO_ROOT)}: {exc}")
            continue

        for node in ast.walk(tree):
            if isinstance(node, ast.Assign):
                for target in node.targets:
                    if (
                        isinstance(target, ast.Name)
                        and target.id == "NODE_CLASS_MAPPINGS"
                    ):
                        rel = py_path.relative_to(REPO_ROOT)
                        raw_keys = _extract_literal_dict_keys(node.value)
                        omitted = raw_keys & _STATIC_NODE_ID_OMISSIONS.get(
                            rel, set()
                        )
                        keys = raw_keys - omitted
                        all_node_ids.update(keys)
                        if verbose and keys:
                            print(_dim(f"  {rel}: {sorted(keys)}"))
                        if verbose and omitted:
                            print(
                                _dim(
                                    f"  {rel}: omitted implementation-only "
                                    f"{sorted(omitted)}"
                                )
                            )

    if parse_errors:
        print(_red(f"  FAIL — {len(parse_errors)} parse errors:"))
        for err in parse_errors:
            print(_red(f"    {err}"))
        return False, all_node_ids

    if not all_node_ids:
        print(_red("  FAIL — no NODE_CLASS_MAPPINGS literal dicts found anywhere"))
        return False, all_node_ids

    print(
        _green(f"  PASS — {files_seen} *.py files parsed, "
               f"{len(all_node_ids)} unique node IDs collected"
               + (":" if verbose else ""))
    )
    if verbose:
        for nid in sorted(all_node_ids):
            print(_dim(f"    {nid}"))
    return True, all_node_ids


# ───────────────────────────────────────────────────────────────────────
# Check 2 — dispatch (stub-VAE roundtrip)
# ───────────────────────────────────────────────────────────────────────


def check_dispatch(verbose: bool = False) -> bool:
    """
    Run the v2_3_3 VAE wrapper through every rank/VAE-type combination using
    a FakeTensor stub. No torch dependency — the stub mimics just enough of
    `torch.Tensor` to drive the dispatch logic.

    Implementation lives in `tools/_preflight_dispatch.py` (extracted to keep
    this orchestrator under the project's ~400-line file-size guideline).
    """
    print(_cyan("dispatch — VAE rank/VAE-type branches"))

    # Lazy import: the sibling module shares the same directory as this file
    # but isn't importable as `tools._preflight_dispatch` unless tools/ is on
    # sys.path. Load it directly via importlib so working dir doesn't matter.
    import importlib.util
    spec = importlib.util.spec_from_file_location(
        "_preflight_dispatch",
        Path(__file__).resolve().parent / "_preflight_dispatch.py",
    )
    dispatch = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(dispatch)

    success, failures = dispatch.run(verbose=verbose, repo_root=REPO_ROOT)

    if failures:
        print(_red(f"  FAIL — {len(failures)} dispatch branch(es) broken:"))
        for f in failures:
            print(_red(f"    {f}"))
        return False
    print(_green("  PASS — all 5 dispatch branches OK"))
    return success


# ───────────────────────────────────────────────────────────────────────
# Check 3 — manager-meta (extension-node-map drift)
# ───────────────────────────────────────────────────────────────────────


def check_manager_meta(actual_ids: set[str], verbose: bool = False) -> bool:
    """
    Compare ComfyUI-Manager's upstream extension-node-map entry for our repo
    against the AST-extracted node IDs. Reports phantom IDs (in upstream but
    not in our code) and missing IDs (in our code but not in upstream).
    """
    print(_cyan("manager-meta — extension-node-map.json drift"))
    if not MANAGER_EXT_MAP_URL.startswith("https://"):
        # Hardcoded URL above is https; this guard ensures any future edit
        # can't accidentally introduce a non-https scheme.
        print(_yellow("  SKIP — manager-meta URL is not https"))
        return True
    try:
        # Custom opener with a redirect handler that re-validates each hop
        # for an https scheme — closes the gap that urlopen by default would
        # follow a 30x to http:// or file:// silently.
        opener = urllib.request.build_opener(_HttpsOnlyRedirectHandler())
        # Bandit B310: scheme is enforced by the guard above + the redirect
        # handler below; the URL is hardcoded to a github.com raw path.
        with opener.open(MANAGER_EXT_MAP_URL, timeout=20) as resp:  # nosec B310
            data = json.loads(resp.read().decode("utf-8"))
    except (urllib.error.URLError, json.JSONDecodeError, TimeoutError) as exc:
        print(_yellow(f"  SKIP — could not fetch upstream map: {exc}"))
        return True

    entry = data.get(REPO_URL)
    if entry is None:
        print(_yellow(f"  SKIP — no entry for {REPO_URL} in upstream map yet"))
        return True

    # Defensive shape-check: upstream entry should be `[list_of_ids, metadata]`.
    # Without this, a future schema change at upstream would silently produce
    # an empty `set()` here and report every one of our IDs as `missing`.
    if not (isinstance(entry, list) and entry and isinstance(entry[0], list)):
        print(_yellow(
            f"  SKIP — upstream entry has unexpected shape: "
            f"{type(entry).__name__}"
        ))
        return True

    upstream_ids = {x for x in entry[0] if isinstance(x, str)}

    phantom = upstream_ids - actual_ids
    missing = actual_ids - upstream_ids

    if phantom or missing:
        print(_red("  FAIL — drift between code and Manager metadata:"))
        if phantom:
            print(_red(f"    Manager lists {len(phantom)} ID(s) not in our code:"))
            for x in sorted(phantom):
                print(_red(f"      - {x}"))
        if missing:
            print(_red(f"    Our code registers {len(missing)} ID(s) not in Manager:"))
            for x in sorted(missing):
                print(_red(f"      - {x}"))
        print(
            _yellow(
                "  -> Manager metadata is community-maintained. "
                "Drift here means an issue or PR to ltdrdata/ComfyUI-Manager. "
                "See issue #44 for the canonical write-up."
            )
        )
        return False

    print(_green(f"  PASS — Manager metadata matches our {len(actual_ids)} IDs"))
    return True


# ───────────────────────────────────────────────────────────────────────
# Check 4 — workflows (fixture node-ID stability)
# ───────────────────────────────────────────────────────────────────────


# Heuristic: a node `type` value in a workflow JSON is treated as a Koolook
# pack ID if it matches any of these patterns.
_KOOLOOK_ID_PATTERNS = [
    re.compile(r"(?i)koolook"),         # contains "koolook" anywhere
    re.compile(r"^Easy[A-Z_]"),         # starts with "Easy" + capital/underscore
    re.compile(r"^easy_[A-Za-z]"),      # starts with "easy_" + any letter
                                        # (covers easy_ImageBatch with capital I)
]


def _looks_like_koolook_id(nid: str) -> bool:
    return any(p.search(nid) for p in _KOOLOOK_ID_PATTERNS)


def _extract_workflow_node_types(data: Any) -> set[str]:
    """
    Extract node `type` / `class_type` strings from a workflow JSON.
    Handles two ComfyUI export formats:
      - Canvas format: top-level dict with a `nodes` array of {type, ...}.
      - API format:    top-level dict where every value is a node spec
                       with a `class_type` key.
    Returns an empty set if the shape doesn't match either format —
    caller decides whether that's a soft or hard failure.
    """
    if not isinstance(data, dict):
        return set()

    # Canvas format
    if isinstance(data.get("nodes"), list):
        return {
            n["type"] for n in data["nodes"]
            if isinstance(n, dict) and isinstance(n.get("type"), str)
        }

    # API format — only accept if EVERY value has class_type, so a malformed
    # canvas (e.g. nodes: null) doesn't accidentally fall through and pull
    # class_type from unrelated keys.
    values = list(data.values())
    if values and all(
        isinstance(v, dict) and isinstance(v.get("class_type"), str) for v in values
    ):
        return {v["class_type"] for v in values}

    return set()


def check_workflows(actual_ids: set[str], verbose: bool = False) -> bool:
    print(_cyan("workflows — fixture node-ID stability"))
    fixtures_dir = REPO_ROOT / "tests" / "workflows"
    if not fixtures_dir.is_dir():
        print(_yellow("  SKIP — no tests/workflows/ folder yet"))
        return True

    fixtures = sorted(fixtures_dir.glob("*.json"))
    if not fixtures:
        # If the dir exists with a README but contains no *.json, that's a
        # regression in fixture coverage (someone deleted them or typoed
        # the dir name). Fail loudly rather than silently pass.
        if (fixtures_dir / "README.md").exists():
            print(_red(
                "  FAIL — fixtures dir exists with README but no *.json files. "
                "If this is intentional, delete the README too."
            ))
            return False
        print(_yellow("  SKIP — no *.json fixtures present"))
        return True

    failures: list[str] = []
    total_koolook_refs = 0

    for fx in fixtures:
        try:
            data = json.loads(fx.read_text(encoding="utf-8"))
        except json.JSONDecodeError as exc:
            failures.append(f"{fx.name}: invalid JSON ({exc})")
            continue

        types = _extract_workflow_node_types(data)

        # Two-pronged matching:
        #   (a) heuristic match — anything that LOOKS like a Koolook ID by
        #       the name pattern. Catches drift where a fixture references
        #       a renamed/removed Koolook ID even if the fork's suffixing
        #       is opaque to the static check.
        #   (b) actual_ids match — anything in the AST-extracted set is by
        #       definition Koolook. Catches IDs that don't fit the heuristic
        #       (e.g. future Koolook nodes named without an "Easy"/"easy_"
        #       prefix and without "koolook" in the ID itself).
        # Union for the "this fixture references N Koolook IDs" report;
        # heuristic-minus-actual for the "missing ID" failure flag.
        koolook_by_heuristic = {t for t in types if _looks_like_koolook_id(t)}
        koolook_by_actual = types & actual_ids
        koolook_total = koolook_by_heuristic | koolook_by_actual
        missing = koolook_by_heuristic - actual_ids

        total_koolook_refs += len(koolook_total)

        if verbose:
            print(_dim(
                f"  {fx.name}: {len(types)} types, "
                f"{len(koolook_total)} Koolook (heuristic={len(koolook_by_heuristic)}, "
                f"actual={len(koolook_by_actual)}), "
                f"{len(missing)} missing"
            ))
        if missing:
            for nid in sorted(missing):
                failures.append(f"{fx.name}: {nid} (no longer registered)")

    if failures:
        print(_red(f"  FAIL — {len(failures)} workflow regression(s):"))
        for f in failures:
            print(_red(f"    {f}"))
        print(
            _yellow(
                "  -> Either add a back-compat alias for the missing ID, "
                "or update the fixture if the rename was intentional + announced."
            )
        )
        return False

    print(_green(
        f"  PASS — {len(fixtures)} fixture(s) reference "
        f"{total_koolook_refs} Koolook node ID(s), all currently registered"
    ))
    return True


# ───────────────────────────────────────────────────────────────────────
# Orchestration
# ───────────────────────────────────────────────────────────────────────


CHECKS = ("static", "dispatch", "manager-meta", "workflows")
DEFAULT_CHECKS = ("static", "dispatch", "workflows")


def main() -> int:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--check", choices=CHECKS, action="append",
        help=(
            "Run only the named check (repeatable). Default: run release gates "
            "(static, dispatch, workflows)."
        ),
    )
    parser.add_argument(
        "--skip", choices=CHECKS, action="append",
        help="Skip the named check (repeatable).",
    )
    parser.add_argument(
        "-v", "--verbose", action="store_true",
        help="Print per-check details.",
    )
    args = parser.parse_args()

    selected = set(args.check or DEFAULT_CHECKS)
    skipped = set(args.skip or ())
    to_run = [c for c in CHECKS if c in selected and c not in skipped]
    if not to_run:
        print(_red("No checks selected after applying --check / --skip"), file=sys.stderr)
        return 2

    print(_cyan(f"Pre-flight running: {', '.join(to_run)}"))
    print()

    actual_ids: set[str] = set()
    results: dict[str, bool] = {}

    if "static" in to_run:
        ok, actual_ids = check_static(verbose=args.verbose)
        results["static"] = ok
        print()

    if "dispatch" in to_run:
        results["dispatch"] = check_dispatch(verbose=args.verbose)
        print()

    if "manager-meta" in to_run:
        if not actual_ids and "static" not in to_run:
            ok, actual_ids = check_static(verbose=False)
        results["manager-meta"] = check_manager_meta(actual_ids, verbose=args.verbose)
        print()

    if "workflows" in to_run:
        if not actual_ids and "static" not in to_run:
            ok, actual_ids = check_static(verbose=False)
        results["workflows"] = check_workflows(actual_ids, verbose=args.verbose)
        print()

    failed = [n for n, ok in results.items() if not ok]

    print(_cyan("Summary"))
    for n in to_run:
        marker = _green("PASS") if results.get(n) else _red("FAIL")
        print(f"  {marker}  {n}")

    if failed:
        print()
        print(_red(f"{len(failed)} check(s) failed: {', '.join(failed)}"))
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())