ausboss-ComfyUI-AusBoss/scripts/validate_nodes.py

#!/usr/bin/env python
"""Offline sanity checks for the node pack — no ComfyUI required.

Run from anywhere:  python scripts/validate_nodes.py

Checks:
  1. Every .py file a registry scanner would read compiles: tracked and
     untracked-but-not-ignored files in a checkout, so local folders such
     as _scratch/ never answer for the pack. Compiled in memory, so the
     check writes no bytecode into the tree it inspects.
  2. Every nodes/node_*.py keeps the registry contract from
     scripts/registry_contract.py: both mappings assigned exactly once, at
     module level, to a non-empty dictionary literal with string-literal
     keys, never mentioned again afterwards, and carrying the same keys as
     each other. Registry scanners parse the source rather than importing
     it, so anything they cannot read statically makes the node invisible.
  3. NODE_MODULES in __init__.py is parsed (not grepped) and matched
     against the files on disk both ways - a listed module with no file and
     a node file nobody imports are both errors - and no mapping key is
     claimed by two modules.
  4. Permanent public node IDs are still registered - measured against the
     mapping keys parsed in check 2, not against the text of the file - and
     each public frame transform node keeps IMAGE, MASK, AUSBOSS_STITCHER,
     IMAGE (image, mask, stitcher, original) as its first outputs; outputs
     appended after them are allowed.
  5. No file outside nodes/ declares mapping keys. Registry scanners read
     the whole checkout, so a fixture or sample that names AUSBOSS_NODES_*
     in a mapping literal is advertised as an installable node.
  6. Every public node's display name ends with the pack suffix " 🆎".
  7. Every public node has its help page at js/docs/<mapping key>.md.

Exit code 0 = all good, 1 = problems printed below.
"""

import ast
import pathlib
import sys

ROOT = pathlib.Path(__file__).resolve().parent.parent
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
from registry_contract import (
    class_mapping_keys,
    duplicate_key_problems,
    mapping_problems,
    module_list,
    registered_class,
    scanned_python_files,
)

errors = []

PUBLIC_NODE_IDS = {
    "AUSBOSS_NODES_AlignImage",
    "AUSBOSS_NODES_Callout",
    "AUSBOSS_NODES_CalloutTile",
    "AUSBOSS_NODES_ColorMatch",
    "AUSBOSS_NODES_Compare",
    "AUSBOSS_NODES_CropForInpaint",
    "AUSBOSS_NODES_Float",
    "AUSBOSS_NODES_FrameInterpolate",
    "AUSBOSS_NODES_FreeMemory",
    "AUSBOSS_NODES_H3PromptWriter",
    "AUSBOSS_NODES_ImageCropRotatePad",
    "AUSBOSS_NODES_ImageFolder",
    "AUSBOSS_NODES_ImageResize",
    "AUSBOSS_NODES_ImageSize",
    "AUSBOSS_NODES_Integer",
    "AUSBOSS_NODES_Krea2Encode",
    "AUSBOSS_NODES_Krea2OutpaintModelPatch",
    "AUSBOSS_NODES_LaMaInpaint",
    "AUSBOSS_NODES_LayoutBoxes",
    "AUSBOSS_NODES_LoadAudio",
    "AUSBOSS_NODES_LoadImagePad",
    "AUSBOSS_NODES_LoadVideo",
    "AUSBOSS_NODES_LoraLoader",
    "AUSBOSS_NODES_MaskByName",
    "AUSBOSS_NODES_MaskByTouch",
    "AUSBOSS_NODES_MathExpression",
    "AUSBOSS_NODES_MergeBatches",
    "AUSBOSS_NODES_RealignToSource",
    "AUSBOSS_NODES_RefineMask",
    "AUSBOSS_NODES_Resolution",
    "AUSBOSS_NODES_RunTimer",
    "AUSBOSS_NODES_SaveImage",
    "AUSBOSS_NODES_SaveVideo",
    "AUSBOSS_NODES_Seed",
    "AUSBOSS_NODES_SelectEveryNth",
    "AUSBOSS_NODES_SelectFrame",
    "AUSBOSS_NODES_ShowText",
    "AUSBOSS_NODES_SplitBatch",
    "AUSBOSS_NODES_StepPreview",
    "AUSBOSS_NODES_StitchInpaint",
    "AUSBOSS_NODES_TiledUpscaleStitch",
    "AUSBOSS_NODES_Text",
    "AUSBOSS_NODES_TiledUpscale",
    "AUSBOSS_NODES_VideoCropRotatePad",
    "AUSBOSS_NODES_VideoCropRotatePadClip",
    "AUSBOSS_NODES_WorkflowNote",
    "AUSBOSS_NODES_WorkflowNoteTile",
    "AUSBOSS_NODES_WorkflowSwitches",
}

# Ids from before the AUSBOSS_NODES_ convention. They stay registered forever
# so workflows already published against them keep loading - the Civitai one
# included. Held in their own set, not folded into PUBLIC_NODE_IDS, so the
# next person to read this list can tell "kept for compatibility" apart from
# "part of the current lineup" and does not tidy one away as a typo.
LEGACY_NODE_IDS = {
    "SimpleWatermarkRemover",
}

RELEASED_NODE_IDS = PUBLIC_NODE_IDS | LEGACY_NODE_IDS

# The pack signature every public display name ends with.
PACK_SUFFIX = " 🆎"

# The permanent first outputs of the two public frame transform nodes.
TRANSFORM_OUTPUT_TYPES = ("IMAGE", "MASK", "AUSBOSS_STITCHER", "IMAGE")
TRANSFORM_OUTPUT_NAMES = ("image", "mask", "stitcher", "original")


def display_names(source):
    """NODE_DISPLAY_NAME_MAPPINGS read the way a scanner reads it: each
    string-literal key with its value, or None where that is not a string
    literal."""
    try:
        tree = ast.parse(source)
    except SyntaxError:
        return {}
    for node in tree.body:
        if (
            isinstance(node, ast.Assign)
            and len(node.targets) == 1
            and isinstance(node.targets[0], ast.Name)
            and node.targets[0].id == "NODE_DISPLAY_NAME_MAPPINGS"
            and isinstance(node.value, ast.Dict)
        ):
            return {
                key.value: value.value if isinstance(value, ast.Constant) else None
                for key, value in zip(node.value.keys, node.value.values)
                if isinstance(key, ast.Constant) and isinstance(key.value, str)
            }
    return {}


def leading(values, count):
    """The first ``count`` entries of a parsed tuple or list, else None."""
    return tuple(values[:count]) if isinstance(values, (tuple, list)) else None


# The files a registry scanner reads: checks 1 and 5 look at nothing else,
# so ignored local folders such as _scratch/ neither fail nor pass the build.
python_files = scanned_python_files(ROOT)

# --- 1. everything compiles --------------------------------------------------
# compile() in memory rather than py_compile, which wrote a __pycache__
# folder next to every file it checked.
for path in python_files:
    label = path.relative_to(ROOT).as_posix()
    try:
        compile(path.read_bytes(), label, "exec", dont_inherit=True)
    except SyntaxError as exc:
        errors.append(f"syntax error: {label}:{exc.lineno}: {exc.msg}")
    except (OSError, ValueError) as exc:
        errors.append(f"cannot compile {label}: {exc}")

# --- 2. the registry contract each node module must keep --------------------
node_files = sorted((ROOT / "nodes").glob("node_*.py"))
keys_by_module = {}
display_by_key = {}
for path in node_files:
    source = path.read_text(encoding="utf-8")
    errors.extend(mapping_problems(source, path.name))
    keys_by_module[path.name] = class_mapping_keys(source)
    display_by_key.update(display_names(source))

# Two modules claiming one key is a silent drop: the last import wins.
errors.extend(duplicate_key_problems(keys_by_module))

# --- 3. NODE_MODULES list matches the files on disk --------------------------
# Parsed out of the list literal, not matched anywhere in the text: a name
# left behind in a comment used to satisfy this, and because an unlisted
# module was only a warning, forgetting to register a node exited 0 with
# nothing printed at all.
init_text = (ROOT / "__init__.py").read_text(encoding="utf-8")
listed = module_list(init_text)
on_disk = {path.stem for path in node_files}

if listed is None:
    errors.append(
        "__init__.py: NODE_MODULES must be a module-level list literal of "
        "module names; this checker reads it without importing the package"
    )
else:
    for missing in sorted(listed - on_disk):
        errors.append(f"__init__.py lists {missing} but nodes/{missing}.py does not exist")
    for orphan in sorted(on_disk - listed):
        # Not a warning: an unlisted module is never imported, so the node
        # simply does not exist in ComfyUI.
        errors.append(
            f"nodes/{orphan}.py is not listed in NODE_MODULES, so it is never imported"
        )

# --- 4. permanent transform contracts ---------------------------------------
# The keys parsed out of NODE_CLASS_MAPPINGS in check 2, not every
# AUSBOSS_NODES_* string in the file: a docstring, a comment or a search
# alias mentioning an id is not a registration, and matching those would let
# a node lose its mapping entry while this check kept passing.
mapping_keys = set().union(*keys_by_module.values()) if keys_by_module else set()
for path in node_files:
    text = path.read_text(encoding="utf-8")
    if path.stem in {"node_image_crop_rotate_pad", "node_video_crop_rotate_pad"}:
        tree = ast.parse(text)
        key = {"node_image_crop_rotate_pad": "AUSBOSS_NODES_ImageCropRotatePad",
               "node_video_crop_rotate_pad": "AUSBOSS_NODES_VideoCropRotatePad"}[path.stem]
        cls = registered_class(tree, key)
        return_types = None
        return_names = None
        for node in cls.body if cls is not None else []:
            if not isinstance(node, ast.Assign):
                continue
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id in {"RETURN_TYPES", "RETURN_NAMES"}:
                    try:
                        value = ast.literal_eval(node.value)
                    except Exception:
                        value = None
                    if target.id == "RETURN_TYPES":
                        return_types = value
                    else:
                        return_names = value
        # A prefix, not the whole tuple. Saved workflows link outputs by slot
        # index, so these four can never move, but an output appended after
        # them leaves every existing link where it was and must not fail the
        # build.
        count = len(TRANSFORM_OUTPUT_TYPES)
        if leading(return_types, count) != TRANSFORM_OUTPUT_TYPES:
            errors.append(f"{path.name}: expected image, mask, stitcher, and original as the first outputs")
        if leading(return_names, count) != TRANSFORM_OUTPUT_NAMES:
            errors.append(f"{path.name}: existing image/mask slots must precede stitcher/original")

for missing_id in sorted(RELEASED_NODE_IDS - mapping_keys):
    errors.append(f"missing permanent mapping key: {missing_id}")

# The same set is a gate in the OTHER direction too. A key nobody listed is a
# node that reached the public pack without anyone deciding it should ship,
# and this catches it. Adding the id above is the deliberate act of
# publishing it, and after a release it can never be renamed: the key is the
# workflow-compatibility contract.
for unlisted_id in sorted(mapping_keys - RELEASED_NODE_IDS):
    errors.append(
        f"unlisted mapping key: {unlisted_id} is registered but not in "
        "PUBLIC_NODE_IDS. Add it there to publish it, or keep it out of the "
        "public pack"
    )

# --- 5. nothing outside nodes/ may declare mapping keys ----------------------
# A scanner reads the whole checkout, not just the modules __init__ imports,
# so any .py anywhere that assigns NODE_CLASS_MAPPINGS advertises those keys
# as installable nodes. Test fixtures did exactly that: they published ids no
# import ever registers, and Manager offered the pack for workflows using
# them. Fixtures live as .py.txt for this reason - real source, not a module.
for path in python_files:
    if path in node_files:
        continue
    stray = class_mapping_keys(path.read_text(encoding="utf-8"))
    for key in sorted(stray):
        errors.append(
            f"{path.relative_to(ROOT)}: declares mapping key {key} outside "
            "a registered node module, where registry scanners will advertise it as a node"
        )

# --- 6. public display names carry the pack suffix ---------------------------
# Names go through ascii(): a Windows console may not encode the emoji, and
# a crash while printing would hide every other problem.
for key in sorted(PUBLIC_NODE_IDS & display_by_key.keys()):
    name = display_by_key[key]
    if not (isinstance(name, str) and name.endswith(PACK_SUFFIX)):
        errors.append(
            f"{key}: display name {ascii(name)} does not end with the pack "
            f"suffix {ascii(PACK_SUFFIX)}"
        )

# --- 7. every public node has a help page -----------------------------------
# The frontend serves js/docs/<mapping key>.md as the node's help panel.
for key in sorted(PUBLIC_NODE_IDS):
    if not (ROOT / "js" / "docs" / f"{key}.md").is_file():
        errors.append(f"{key} has no help page: add js/docs/{key}.md")

# --- report ------------------------------------------------------------------
for error in errors:
    print(f"ERROR {error}")

if errors:
    print(f"\n{len(errors)} problem(s) found.")
    sys.exit(1)

print(f"OK - {len(node_files)} node file(s) validated.")