Merge pull request #3 from relhamdi/node-prompt_switch

PromptSwitch

README.md

@@ -169,7 +169,6 @@ Multiple `@combine` blocks are allowed and processed independently. A `---` outs
 #### Validation
 
 The node stops the workflow on error if:
-- `options` is empty
 - A `@combine` block is opened inside another
 - A `@end` is found without a matching `@combine`
 - A `@combine` block is never closed
@@ -229,6 +228,36 @@ Slot lengths do not need to match.
 
 ---
 
+### PromptSwitch
+
+![PromptSwitch_v1](docs/images/PromptSwitch_v1.png)
+
+#### Inputs
+
+- `on_true` (str) — value returned when condition is `True`. Accepts direct input or connected node.
+- `on_false` (str) — value returned when condition is `False`. Accepts direct input or connected node.
+- `condition` (bool), default `True` — controls which input is routed to the output.
+
+#### Outputs
+
+- `text` (str) — the selected value based on condition.
+
+#### Usage
+
+Connect or type a value in `on_true` and `on_false`. Toggle `condition` to switch between them.
+
+If the node is bypassed, an empty string is returned.
+
+#### Editor
+
+Input ports and widget borders are colored based on the active condition:
+- **Orange**: active input when condition is `True`
+- **Blue**: active input when condition is `False`
+- **Gray**: inactive input
+Colors are only shown when an input is connected or has a value.
+
+---
+
 ### Web
 
 #### contentEditable editor
@@ -268,6 +297,8 @@ The selection is preserved after each keypress, consistent with ComfyUI native b
 
 ## TODO
 
+- Centralize colors used
+
 `PromptPresetSelector`
 - Add a field to easily concatenate text at the end of the output
-- Wildcard mode for presets (randomly selected, index to -1 or -2?)
+- Wildcard mode for presets (randomly selected, index to -1 or -2?)

__init__.py

@@ -18,16 +18,24 @@ from .src.nodes.prompt_preset_selector import (
 from .src.nodes.prompt_preset_selector import (
     NODE_DISPLAY_NAME_MAPPINGS as PromptPresetSelectorDisplayMappings,
 )
+from .src.nodes.prompt_switch import (
+    NODE_CLASS_MAPPINGS as PromptSwitchMappings,
+)
+from .src.nodes.prompt_switch import (
+    NODE_DISPLAY_NAME_MAPPINGS as PromptSwitchDisplayMappings,
+)
 
 NODE_CLASS_MAPPINGS = {}
 NODE_CLASS_MAPPINGS.update(PromptPresetSelectorMappings)
 NODE_CLASS_MAPPINGS.update(PromptOptionPickerMappings)
 NODE_CLASS_MAPPINGS.update(PromptLayoutFillerMappings)
+NODE_CLASS_MAPPINGS.update(PromptSwitchMappings)
 
 NODE_DISPLAY_NAME_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS.update(PromptPresetSelectorDisplayMappings)
 NODE_DISPLAY_NAME_MAPPINGS.update(PromptOptionPickerDisplayMappings)
 NODE_DISPLAY_NAME_MAPPINGS.update(PromptLayoutFillerDisplayMappings)
+NODE_DISPLAY_NAME_MAPPINGS.update(PromptSwitchDisplayMappings)
 
 # Node registration
 __all__ = [

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ComfyUI-RE-CustomUtils"
-version = "0.3.0"
+version = "0.4.0"
 description = "A collection of custom nodes for ComfyUI."
 license = { file = "LICENSE" }
 readme = "README.md"

src/nodes/prompt_option_picker.py

@@ -42,14 +42,13 @@ class PromptOptionPicker:
 
     @classmethod
     def VALIDATE_INPUTS(cls, options, **kwargs):
-        if not options or not options.strip():
-            return "options cannot be empty"
+        try:
+            parse_options(options)
+        except ValueError as e:
+            return str(e)
         return True
 
     def process(self, options, selected):
-        # Validate only, JS handles the dropdown
-        parse_options(options)
-
         # -- displays an empty string
         return ("" if selected == "--" else selected,)
 

src/nodes/prompt_switch.py

@@ -0,0 +1,53 @@
+from .. import config
+
+
+class PromptSwitch:
+    """
+    Routes one of two STRING inputs to the output based on a boolean condition.
+    If condition is True, returns on_true. If False, returns on_false.
+    If the node is bypassed, ComfyUI returns on_true by convention.
+    """
+
+    CATEGORY = config.NODE_CATEGORY
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "on_true": (
+                    "STRING",
+                    {
+                        "tooltip": "Value returned when condition is True.",
+                    },
+                ),
+                "on_false": (
+                    "STRING",
+                    {
+                        "tooltip": "Value returned when condition is False.",
+                    },
+                ),
+                "condition": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": "If True, returns on_true. If False, returns on_false.",
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("text",)
+    FUNCTION = "process"
+
+    def process(self, on_true: str, on_false: str, condition: bool) -> tuple[str]:
+        return (on_true if condition else on_false,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "PromptSwitch": PromptSwitch,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PromptSwitch": "Prompt Switch",
+}

tests/test_prompt_option_picker.py

@@ -41,17 +41,22 @@ def test_whitespace_only_lines(node):
     assert value == ""
 
 
+def test_empty_options_returns_empty(node):
+    (value,) = node.process(options="", selected="--")
+    assert value == ""
+
+
 # --- VALIDATE_INPUTS ---
 
 
 def test_validate_empty_options():
     result = PromptOptionPicker.VALIDATE_INPUTS(options="")
-    assert result != True
+    assert result == True
 
 
 def test_validate_whitespace_only():
     result = PromptOptionPicker.VALIDATE_INPUTS(options="   ")
-    assert result != True
+    assert result == True
 
 
 def test_validate_valid():

tests/test_prompt_switch.py

@@ -0,0 +1,67 @@
+import pytest
+from src.nodes.prompt_switch import PromptSwitch
+
+
+@pytest.fixture
+def node():
+    return PromptSwitch()
+
+
+def run(node, on_true="", on_false="", condition=True):
+    return node.process(on_true=on_true, on_false=on_false, condition=condition)
+
+
+# --- Base cases ---
+
+
+def test_condition_true(node):
+    (text,) = run(node, on_true="hello", on_false="world", condition=True)
+    assert text == "hello"
+
+
+def test_condition_false(node):
+    (text,) = run(node, on_true="hello", on_false="world", condition=False)
+    assert text == "world"
+
+
+def test_empty_on_true(node):
+    (text,) = run(node, on_true="", on_false="world", condition=True)
+    assert text == ""
+
+
+def test_empty_on_false(node):
+    (text,) = run(node, on_true="hello", on_false="", condition=False)
+    assert text == ""
+
+
+def test_both_empty(node):
+    (text,) = run(node, on_true="", on_false="", condition=True)
+    assert text == ""
+
+
+def test_multiline_value(node):
+    (text,) = run(node, on_true="line1\nline2", on_false="other", condition=True)
+    assert text == "line1\nline2"
+
+
+# --- With connected nodes (values passed as strings) ---
+
+
+def test_connected_node_true(node):
+    (text,) = run(
+        node,
+        on_true="masterpiece, detailed",
+        on_false="simple",
+        condition=True,
+    )
+    assert text == "masterpiece, detailed"
+
+
+def test_connected_node_false(node):
+    (text,) = run(
+        node,
+        on_true="detailed",
+        on_false="{% preset_a | preset_b %}",
+        condition=False,
+    )
+    assert text == "{% preset_a | preset_b %}"

web/prompt_layout_filler.js

@@ -118,7 +118,7 @@ const attachEditor = (node) => {
         renderColored();
     };
 
-    // --- Initial render ---
+    // Initial render
     updateSlotVisibility(node);
     renderColored();
 };

web/prompt_switch.js

@@ -0,0 +1,127 @@
+import { hookWidget } from "./utils.js";
+import { app } from "/scripts/app.js";
+
+// --- Constants ---
+
+const NODE_NAME = "PromptSwitch";
+
+const COLORS = {
+    true: "#ff9800",
+    false: "#64b5f6",
+    inactive: "#888888",
+};
+
+// --- Highlight ---
+
+const updateInputColors = (node, condition) => {
+    const trueInput = node.inputs?.find((inp) => inp.name === "on_true");
+    const falseInput = node.inputs?.find((inp) => inp.name === "on_false");
+    const trueWidget = node.widgets?.find((w) => w.name === "on_true");
+    const falseWidget = node.widgets?.find((w) => w.name === "on_false");
+    if (!trueInput || !falseInput) return;
+
+    const trueActive =
+        trueInput.link != null || (trueWidget?.value ?? "").trim() !== "";
+    const falseActive =
+        falseInput.link != null || (falseWidget?.value ?? "").trim() !== "";
+
+    trueInput.color_on = trueActive
+        ? condition
+            ? COLORS.true
+            : COLORS.inactive
+        : undefined;
+    trueInput.color_off = trueInput.color_on;
+
+    falseInput.color_on = falseActive
+        ? condition
+            ? COLORS.inactive
+            : COLORS.false
+        : undefined;
+    falseInput.color_off = falseInput.color_on;
+
+    if (node.graph) node.graph.setDirtyCanvas(true, true);
+};
+
+// --- Attach ---
+
+const attachSwitch = (node) => {
+    if (node.type !== NODE_NAME) return;
+    if (node._switchAttached) return;
+    node._switchAttached = true;
+
+    const trueWidget = node.widgets?.find((w) => w.name === "on_true");
+    const falseWidget = node.widgets?.find((w) => w.name === "on_false");
+    const conditionWidget = node.widgets?.find((w) => w.name === "condition");
+    if (!conditionWidget) return;
+
+    const refresh = () => {
+        const condition = conditionWidget.value ?? true;
+        updateInputColors(node, condition);
+        if (node.graph) node.graph.setDirtyCanvas(true, true);
+    };
+
+    // Input port colors
+    hookWidget(conditionWidget, refresh);
+    hookWidget(trueWidget, refresh);
+    hookWidget(falseWidget, refresh);
+
+    // Widget border via canvas
+    const original = node.onDrawForeground;
+    node.onDrawForeground = function (ctx) {
+        if (original) original.call(this, ctx);
+
+        const trueWidget = node.widgets?.find((w) => w.name === "on_true");
+        const falseWidget = node.widgets?.find((w) => w.name === "on_false");
+        const trueInput = node.inputs?.find((inp) => inp.name === "on_true");
+        const falseInput = node.inputs?.find((inp) => inp.name === "on_false");
+        if (!trueWidget || !falseWidget) return;
+
+        const condition =
+            node.widgets?.find((w) => w.name === "condition")?.value ?? true;
+
+        const drawBorder = (widget, color, input) => {
+            // Skip if not connected AND widget value is empty
+            const isConnected = input?.link != null;
+            const hasValue = widget.value != null && widget.value.trim() !== "";
+            if (!isConnected && !hasValue) return;
+
+            const x = 0;
+            const y = widget.last_y - 2;
+            const w = node.size[0];
+            const h = widget.computedHeight ?? 20;
+
+            ctx.strokeStyle = color;
+            ctx.lineWidth = 2;
+            ctx.strokeRect(x, y, w, h);
+        };
+
+        // Draw inactive first, active on top
+        if (condition) {
+            drawBorder(falseWidget, COLORS.inactive, falseInput);
+            drawBorder(trueWidget, COLORS.true, trueInput);
+        } else {
+            drawBorder(trueWidget, COLORS.inactive, trueInput);
+            drawBorder(falseWidget, COLORS.false, falseInput);
+        }
+    };
+
+    // Initial render
+    refresh();
+};
+
+// --- Registration ---
+
+app.registerExtension({
+    name: NODE_NAME,
+    beforeRegisterNodeDef(nodeType, nodeData) {
+        if (nodeData.name !== NODE_NAME) return;
+
+        const original = nodeType.prototype.onNodeCreated;
+        nodeType.prototype.onNodeCreated = function () {
+            if (original) original.call(this);
+
+            const node = this;
+            requestAnimationFrame(() => attachSwitch(node));
+        };
+    },
+});