From f5908006ebb2db80d704f8d949df1f09c7cdd4c2 Mon Sep 17 00:00:00 2001 From: Crespin Alexandre <114501859+kunail0804@users.noreply.github.com> Date: Wed, 8 Jul 2026 14:16:42 +0200 Subject: [PATCH] feat(nodes): add the Category Splitter ComfyUI node (#22) (#66) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add nodes/category_splitter_node.py: CategorySplitterNode, a thin adapter that delegates to split_into_categories (MASTER_SPEC §20). - Input resolved_tags (RESOLVED_TAGS); outputs (category_map, warnings, errors): the CATEGORY_MAP (per-category grouping, correct assignment + order) plus surfaced diagnostics. No compiler logic in the node. - The per-category prompt strings are produced downstream by the Prompt Builder node (§19); this node produces the intermediate Category Map that feeds it. - Registered as SceneCompilerCategorySplitter ("Category Splitter"). Tests: metadata/inputs, real split into a CategoryMap, delegation + SC0008 surfacing (monkeypatched), and registration. Co-authored-by: kunail0804 Co-authored-by: Claude Opus 4.8 --- nodes/__init__.py | 3 ++ nodes/category_splitter_node.py | 35 ++++++++++++++++++++++ tests/test_category_splitter_node.py | 45 ++++++++++++++++++++++++++++ 3 files changed, 83 insertions(+) create mode 100644 nodes/category_splitter_node.py create mode 100644 tests/test_category_splitter_node.py 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"