ComfyUI-Ethanfel-Prompt-Builder/tools/prompt_map_audit.py

#!/usr/bin/env python3
"""Print a lightweight audit for the prompt routing map.

This intentionally avoids importing the ComfyUI node package. It parses Python
and JSON files directly, so it can run in a plain shell without ComfyUI loaded.
"""

from __future__ import annotations

import ast
import json
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]


def _literal_or_none(node: ast.AST) -> Any:
    try:
        return ast.literal_eval(node)
    except Exception:
        return None


def _assignment_dict(path: Path, name: str) -> dict[str, Any]:
    tree = ast.parse(path.read_text(encoding="utf-8"))
    for node in tree.body:
        if not isinstance(node, ast.Assign):
            continue
        if not any(isinstance(target, ast.Name) and target.id == name for target in node.targets):
            continue
        value = _literal_or_none(node.value)
        return value if isinstance(value, dict) else {}
    return {}


def _class_return_names(path: Path) -> dict[str, tuple[str, ...]]:
    tree = ast.parse(path.read_text(encoding="utf-8"))
    result: dict[str, tuple[str, ...]] = {}
    for node in tree.body:
        if not isinstance(node, ast.ClassDef) or not node.name.startswith("SxCP"):
            continue
        for item in node.body:
            if not isinstance(item, ast.Assign):
                continue
            if not any(isinstance(target, ast.Name) and target.id == "RETURN_NAMES" for target in item.targets):
                continue
            value = _literal_or_none(item.value)
            if isinstance(value, tuple) and all(isinstance(part, str) for part in value):
                result[node.name] = value
    return result


def _category_summary(path: Path) -> dict[str, Any]:
    data = json.loads(path.read_text(encoding="utf-8"))
    categories = data.get("categories") or []
    subcategory_count = 0
    item_template_count = 0
    for category in categories:
        subcategories = category.get("subcategories") or []
        subcategory_count += len(subcategories)
        for subcategory in subcategories:
            item_template_count += len(subcategory.get("item_templates") or [])
            for item in subcategory.get("items") or []:
                if isinstance(item, dict):
                    item_template_count += len(item.get("item_templates") or [])
    return {
        "categories": len(categories),
        "subcategories": subcategory_count,
        "item_templates": item_template_count,
        "scene_pools": len(data.get("scene_pools") or {}),
        "expression_pools": len(data.get("expression_pools") or {}),
        "composition_pools": len(data.get("composition_pools") or {}),
        "pool_extensions": len(data.get("pool_extensions") or {}),
    }


def _pool_names(path: Path, key: str) -> list[str]:
    data = json.loads(path.read_text(encoding="utf-8"))
    pools = data.get(key) or {}
    return sorted(pools) if isinstance(pools, dict) else []


def print_table(headers: tuple[str, ...], rows: list[tuple[Any, ...]]) -> None:
    widths = [len(header) for header in headers]
    for row in rows:
        for index, value in enumerate(row):
            widths[index] = max(widths[index], len(str(value)))
    print("| " + " | ".join(header.ljust(widths[index]) for index, header in enumerate(headers)) + " |")
    print("| " + " | ".join("-" * width for width in widths) + " |")
    for row in rows:
        print("| " + " | ".join(str(value).ljust(widths[index]) for index, value in enumerate(row)) + " |")


def main() -> int:
    init_path = ROOT / "__init__.py"
    loop_path = ROOT / "loop_nodes.py"
    display = _assignment_dict(init_path, "NODE_DISPLAY_NAME_MAPPINGS")
    loop_display = _assignment_dict(loop_path, "LOOP_NODE_DISPLAY_NAME_MAPPINGS")
    display.update(loop_display)
    returns = _class_return_names(init_path)
    returns.update(_class_return_names(loop_path))

    print("# Node Display Map")
    node_rows = []
    for class_name, display_name in sorted(display.items(), key=lambda item: str(item[1])):
        return_names = ", ".join(returns.get(class_name, ()))
        node_rows.append((display_name, class_name, return_names or "(dynamic or unnamed)"))
    print_table(("Display name", "Class", "Return names"), node_rows)

    print("\n# Category JSON Summary")
    category_rows = []
    for path in sorted((ROOT / "categories").glob("*.json")):
        summary = _category_summary(path)
        category_rows.append(
            (
                path.name,
                summary["categories"],
                summary["subcategories"],
                summary["item_templates"],
                summary["scene_pools"],
                summary["expression_pools"],
                summary["composition_pools"],
                summary["pool_extensions"],
            )
        )
    print_table(
        (
            "File",
            "Categories",
            "Subcategories",
            "Item templates",
            "Scene pools",
            "Expression pools",
            "Composition pools",
            "Extensions",
        ),
        category_rows,
    )

    print("\n# Named Pool Inventory")
    pool_rows = []
    for path in sorted((ROOT / "categories").glob("*.json")):
        for key in ("scene_pools", "expression_pools", "composition_pools"):
            names = _pool_names(path, key)
            if names:
                pool_rows.append((path.name, key, len(names), ", ".join(names[:8]) + (" ..." if len(names) > 8 else "")))
    print_table(("File", "Pool type", "Count", "First names"), pool_rows)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())