diff --git a/nodes/__init__.py b/nodes/__init__.py index 23a2be0..ff17a60 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -7,6 +7,7 @@ Nodes are thin interfaces; all compiler logic lives in the ``compiler`` package. from __future__ import annotations +from .category_splitter_node import CategorySplitterNode from .resolver_node import ResolverNode from .scene_analyzer_node import SceneAnalyzerNode from .scene_validator_node import SceneValidatorNode @@ -15,12 +16,14 @@ NODE_CLASS_MAPPINGS: dict[str, type] = { "SceneCompilerAnalyzer": SceneAnalyzerNode, "SceneCompilerValidator": SceneValidatorNode, "SceneCompilerResolver": ResolverNode, + "SceneCompilerCategorySplitter": CategorySplitterNode, } NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = { "SceneCompilerAnalyzer": "Scene Analyzer", "SceneCompilerValidator": "Scene Validator", "SceneCompilerResolver": "Resolver", + "SceneCompilerCategorySplitter": "Category Splitter", } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/category_splitter_node.py b/nodes/category_splitter_node.py new file mode 100644 index 0000000..2e184a5 --- /dev/null +++ b/nodes/category_splitter_node.py @@ -0,0 +1,35 @@ +"""ComfyUI node for the Category Splitter (MASTER_SPEC §20). + +A thin interface: it adapts the Resolved Tags input to the Category Splitter +module and surfaces the CompilerResult's data, warnings, and errors. No compiler +logic; it never imports ComfyUI. + +The node outputs the Category Map (per-category grouping, in order); the +per-category prompt strings are produced downstream by the Prompt Builder node +(§19). +""" + +from __future__ import annotations + +from typing import Any + +from compiler.splitter.category_splitter import split_into_categories + +from .adapters import format_messages + + +class CategorySplitterNode: + """Groups Resolved Tags into a Category Map.""" + + CATEGORY = "Scene Compiler" + FUNCTION = "run" + RETURN_TYPES = ("CATEGORY_MAP", "STRING", "STRING") + RETURN_NAMES = ("category_map", "warnings", "errors") + + @classmethod + def INPUT_TYPES(cls) -> dict[str, Any]: + return {"required": {"resolved_tags": ("RESOLVED_TAGS",)}} + + def run(self, resolved_tags: Any) -> tuple[Any, str, str]: + result = split_into_categories(tuple(resolved_tags)) + return (result.data, format_messages(result.warnings), format_messages(result.errors)) diff --git a/tests/test_category_splitter_node.py b/tests/test_category_splitter_node.py new file mode 100644 index 0000000..5916268 --- /dev/null +++ b/tests/test_category_splitter_node.py @@ -0,0 +1,45 @@ +"""Tests for the Category Splitter ComfyUI node (issue #22).""" + +from __future__ import annotations + +import nodes +import nodes.category_splitter_node as node_module +from compiler.common.message_codes import message +from compiler.common.result import CompilerResult +from nodes.category_splitter_node import CategorySplitterNode +from schemas.models import CategoryMap, ResolvedTag + + +def rtag(name: str, category: str) -> ResolvedTag: + return ResolvedTag(tag=name, category=category, source_concept=name, knowledge_base_entry=name) + + +def test_node_metadata() -> None: + assert CategorySplitterNode.RETURN_TYPES == ("CATEGORY_MAP", "STRING", "STRING") + assert CategorySplitterNode.RETURN_NAMES == ("category_map", "warnings", "errors") + assert "resolved_tags" in CategorySplitterNode.INPUT_TYPES()["required"] + + +def test_splits_with_real_module() -> None: + tags = (rtag("1girl", "character"), rtag("long hair", "hair")) + category_map, warnings, errors = CategorySplitterNode().run(tags) + assert isinstance(category_map, CategoryMap) + assert [t.tag for t in category_map.tags_for("character")] == ["1girl"] + assert [t.tag for t in category_map.tags_for("hair")] == ["long hair"] + assert (warnings, errors) == ("", "") + + +def test_delegates_and_surfaces_errors(monkeypatch) -> None: + monkeypatch.setattr( + node_module, + "split_into_categories", + lambda tags: CompilerResult().add_error(message("SC0008", "bad category")), + ) + category_map, warnings, errors = CategorySplitterNode().run(()) + assert category_map is None + assert "SC0008: bad category" in errors + + +def test_registered_for_comfyui() -> None: + assert nodes.NODE_CLASS_MAPPINGS["SceneCompilerCategorySplitter"] is CategorySplitterNode + assert nodes.NODE_DISPLAY_NAME_MAPPINGS["SceneCompilerCategorySplitter"] == "Category Splitter"