Add rage nodes, disable flow and force calculation nodes

2025-06-09 13:33:10 +02:00
parent 1296aed6a5
commit 622f471bfa
5 changed files with 427 additions and 17 deletions
@@ -90,7 +90,7 @@ class IfElifElse(ComfyNodeABC):
        # Check if condition
        if kwargs.get("if", False) and kwargs.get("then") is None:
            needed.append("then")
-            return needed  # If main condition is true, we only need "then"
+            return needed  # If the main condition is true, we only need "then"

        # Check each elif condition
        elif_index = 0
@@ -153,7 +153,7 @@ class SwitchCase(ComfyNodeABC):
    def INPUT_TYPES(cls):
        return {
            "required": ContainsDynamicDict({
-                "selector": (IO.INT, {"default": 0, "min": 0}),
+                "select": (IO.INT, {"default": 0, "min": 0}),
                "case_0": (IO.ANY, {"lazy": True, "_dynamic": "number"}),
            }),
            "optional": {
@@ -167,29 +167,29 @@ class SwitchCase(ComfyNodeABC):
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "execute"

-    def check_lazy_status(self, selector: int, **kwargs) -> list[str]:
+    def check_lazy_status(self, select: int, **kwargs) -> list[str]:
        needed = []

-        # Check for needed case inputs based on selector
+        # Check for necessary case inputs based on select
        case_count = 0
        for key, value in kwargs.items():
            if key.startswith("case_"):
                try:
                    case_index = int(key.split("_")[1])
                    case_count = max(case_count, case_index + 1)
-                    if value is None and selector == case_index:
+                    if value is None and select == case_index:
                        needed.append(key)
                except ValueError:
                    pass  # Not a numeric case key

-        # Check if default is needed when selector is out of range
-        if "default" in kwargs and kwargs["default"] is None and not 0 <= selector < case_count:
+        # Check if default is needed when select is out of range
+        if "default" in kwargs and kwargs["default"] is None and not 0 <= select < case_count:
            needed.append("default")

        return needed

    def execute(self, selector: int, **kwargs) -> tuple[Any]:
-        # Build cases array from all case_X inputs
+        # Build a case array from all case_X inputs
        cases = []
        for i in range(len(kwargs)):
            case_key = f"case_{i}"
@@ -202,10 +202,44 @@ class SwitchCase(ComfyNodeABC):
        if 0 <= selector < len(cases) and cases[selector] is not None:
            return (cases[selector],)

-        # If selector is out of range or the selected case is None, return default
+        # If select is out of range or the selected case is None, return default
        return (kwargs.get("default"),)


+class DisableFlow(ComfyNodeABC):
+    """
+    Conditionally enable or disable a flow.
+
+    This node takes a value and either passes it through or blocks execution
+    based on the 'select' parameter. When 'select' is True, the value passes through;
+    when False, execution is blocked.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "value": (IO.ANY, {}),
+                "select": (IO.BOOLEAN, {"default": True}),
+            }
+        }
+
+    RETURN_TYPES = (IO.ANY,)
+    RETURN_NAMES = ("value",)
+    CATEGORY = "Basic/flow control"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "execute"
+
+    @classmethod
+    def IS_CHANGED(s, value: Any):
+        return float("NaN") # not equal to anything -> trigger recalculation
+
+    def execute(self, value: Any, select: bool = True) -> tuple[Any]:
+        if select:
+            return (value,)
+        else:
+            return (ExecutionBlocker(None),)
+
+
 class FlowSelect(ComfyNodeABC):
    """
    Select the direction of the flow.
@@ -237,6 +271,35 @@ class FlowSelect(ComfyNodeABC):
            return ExecutionBlocker(None), value


+class ForceCalculation(ComfyNodeABC):
+    """
+    Forces recalculation of the connected nodes.
+
+    This node passes the input directly to the output but prevents caching
+    by marking itself as an output node and also indicates the out has changed.
+    Use this when you need to ensure nodes are always recalculated.
+    """
+
+    OUTPUT_NODE = True  # Marks as an output node to force calculation
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "value": (IO.ANY, {}),
+            }
+        }
+
+    RETURN_TYPES = (IO.ANY,)
+    RETURN_NAMES = ("value",)
+    CATEGORY = "Basic/flow control"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "execute"
+
+    def execute(self, value: Any) -> tuple[Any, int]:
+        return (value,)
+
+
 class ExecutionOrder(ComfyNodeABC):
    """
    Force execution order in the workflow.
@@ -268,7 +331,9 @@ NODE_CLASS_MAPPINGS = {
    "Basic data handling: IfElse": IfElse,
    "Basic data handling: IfElifElse": IfElifElse,
    "Basic data handling: SwitchCase": SwitchCase,
+    "Basic data handling: DisableFlow": DisableFlow,
    "Basic data handling: FlowSelect": FlowSelect,
+    "Basic data handling: ForceCalculation": ForceCalculation,
    "Basic data handling: ExecutionOrder": ExecutionOrder,
 }

@@ -276,6 +341,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "Basic data handling: IfElse": "if/else",
    "Basic data handling: IfElifElse": "if/elif/.../else",
    "Basic data handling: SwitchCase": "switch/case",
+    "Basic data handling: DisableFlow": "disable flow",
    "Basic data handling: FlowSelect": "flow select",
+    "Basic data handling: ForceCalculation": "force calculation",
    "Basic data handling: ExecutionOrder": "force execution order",
 }
@@ -675,6 +675,37 @@ class DataListPopRandom(ComfyNodeABC):
        return input_list, None


+class DataListRange(ComfyNodeABC):
+    """
+    Creates a data list containing a sequence of numbers.
+
+    This node generates a sequence of numbers similar to Python's range() function.
+    It takes start, stop, and step parameters to define the sequence.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "start": ("INT", {"default": 0}),
+                "stop": ("INT", {"default": 10}),
+            },
+            "optional": {
+                "step": ("INT", {"default": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "create_range"
+    CATEGORY = "Basic/Data List"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    OUTPUT_IS_LIST = (True,)
+
+    def create_range(self, stop: int, start: int = 0, step: int = 1) -> tuple[list[int]]:
+        if step == 0:
+            raise ValueError("Step cannot be zero")
+        return (list(range(start, stop, step)),)
+
+
 class DataListRemove(ComfyNodeABC):
    """
    Removes the first occurrence of a specified value from a list.
@@ -960,6 +991,7 @@ NODE_CLASS_MAPPINGS = {
    "Basic data handling: DataListMin": DataListMin,
    "Basic data handling: DataListPop": DataListPop,
    "Basic data handling: DataListPopRandom": DataListPopRandom,
+    "Basic data handling: DataListRange": DataListRange,
    "Basic data handling: DataListRemove": DataListRemove,
    "Basic data handling: DataListReverse": DataListReverse,
    "Basic data handling: DataListSetItem": DataListSetItem,
@@ -992,6 +1024,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "Basic data handling: DataListMin": "min",
    "Basic data handling: DataListPop": "pop",
    "Basic data handling: DataListPopRandom": "pop random",
+    "Basic data handling: DataListRange": "range",
    "Basic data handling: DataListRemove": "remove",
    "Basic data handling: DataListReverse": "reverse",
    "Basic data handling: DataListSetItem": "set item",
@@ -548,6 +548,36 @@ class ListPopRandom(ComfyNodeABC):
        return result, None


+class ListRange(ComfyNodeABC):
+    """
+    Creates a LIST containing a sequence of numbers.
+
+    This node generates a LIST of numbers similar to Python's range() function.
+    It takes start, stop, and step parameters to define the sequence.
+    """
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "start": ("INT", {"default": 0}),
+                "stop": ("INT", {"default": 10}),
+            },
+            "optional": {
+                "step": ("INT", {"default": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("LIST",)
+    CATEGORY = "Basic/LIST"
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "create_range"
+
+    def create_range(self, start: int, stop: int, step: int = 1) -> tuple[list[int]]:
+        if step == 0:
+            raise ValueError("Step cannot be zero")
+        return (list(range(start, stop, step)),)
+
+
 class ListRemove(ComfyNodeABC):
    """
    Removes the first occurrence of a specified value from a LIST.
@@ -727,7 +757,6 @@ class ListToDataList(ComfyNodeABC):
    def convert(self, list) -> tuple[list[Any]]:
        return (list,)

-
 class ListToSet(ComfyNodeABC):
    """
    Converts a LIST into a SET.
@@ -772,6 +801,7 @@ NODE_CLASS_MAPPINGS = {
    "Basic data handling: ListMin": ListMin,
    "Basic data handling: ListPop": ListPop,
    "Basic data handling: ListPopRandom": ListPopRandom,
+    "Basic data handling: ListRange": ListRange,
    "Basic data handling: ListRemove": ListRemove,
    "Basic data handling: ListReverse": ListReverse,
    "Basic data handling: ListSetItem": ListSetItem,
@@ -801,6 +831,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    "Basic data handling: ListMin": "min",
    "Basic data handling: ListPop": "pop",
    "Basic data handling: ListPopRandom": "pop random",
+    "Basic data handling: ListRange": "range",
    "Basic data handling: ListRemove": "remove",
    "Basic data handling: ListReverse": "reverse",
    "Basic data handling: ListSetItem": "set item",
@@ -1,13 +1,14 @@
 import pytest
 from src.basic_data_handling.data_list_nodes import (
    DataListAppend,
-    #DataListCreate,
-    #DataListCreateFromBoolean,
-    #DataListCreateFromFloat,
-    #DataListCreateFromInt,
-    #DataListCreateFromString,
+    DataListCreate,
+    DataListCreateFromBoolean,
+    DataListCreateFromFloat,
+    DataListCreateFromInt,
+    DataListCreateFromString,
    DataListExtend,
    DataListInsert,
+    DataListRange,
    DataListRemove,
    DataListPop,
    DataListPopRandom,
@@ -27,8 +28,8 @@ from src.basic_data_handling.data_list_nodes import (
    DataListLast,
    DataListMin,
    DataListMax,
-    #DataListToList,
-    #DataListToSet,
+    DataListToList,
+    DataListToSet,
 )


@@ -174,6 +175,186 @@ def test_filter_select():
    assert false_list == [1, 2, 3]


+def test_create():
+    node = DataListCreate()
+    # Testing with one item
+    assert node.create_list(item_0="test", _dynamic_number=1) == (["test"],)
+
+    # Testing with multiple items of different types
+    assert node.create_list(item_0=1, item_1="two", item_2=3.0, _dynamic_number=3) == ([1, "two", 3.0],)
+
+    # Testing with empty list (no items)
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_boolean():
+    node = DataListCreateFromBoolean()
+    # Testing with boolean values
+    assert node.create_list(item_0=True, item_1=False, _dynamic_number=2) == ([True, False],)
+
+    # Testing with boolean-convertible values
+    assert node.create_list(item_0=1, item_1=0, _dynamic_number=2) == ([True, False],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_float():
+    node = DataListCreateFromFloat()
+    # Testing with float values
+    assert node.create_list(item_0=1.5, item_1=2.5, _dynamic_number=2) == ([1.5, 2.5],)
+
+    # Testing with float-convertible values
+    assert node.create_list(item_0=1, item_1="2.5", _dynamic_number=2) == ([1.0, 2.5],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_int():
+    node = DataListCreateFromInt()
+    # Testing with integer values
+    assert node.create_list(item_0=1, item_1=2, _dynamic_number=2) == ([1, 2],)
+
+    # Testing with int-convertible values
+    assert node.create_list(item_0="1", item_1=2.0, _dynamic_number=2) == ([1, 2],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_string():
+    node = DataListCreateFromString()
+    # Testing with string values
+    assert node.create_list(item_0="hello", item_1="world", _dynamic_number=2) == (["hello", "world"],)
+
+    # Testing with string-convertible values
+    assert node.create_list(item_0=123, item_1=True, _dynamic_number=2) == (["123", "True"],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create():
+    node = DataListCreate()
+    # Testing with one item
+    assert node.create_list(item_0="test", _dynamic_number=1) == (["test"],)
+
+    # Testing with multiple items of different types
+    assert node.create_list(item_0=1, item_1="two", item_2=3.0, _dynamic_number=3) == ([1, "two", 3.0],)
+
+    # Testing with empty list (no items)
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_boolean():
+    node = DataListCreateFromBoolean()
+    # Testing with boolean values
+    assert node.create_list(item_0=True, item_1=False, _dynamic_number=2) == ([True, False],)
+
+    # Testing with boolean-convertible values
+    assert node.create_list(item_0=1, item_1=0, _dynamic_number=2) == ([True, False],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_float():
+    node = DataListCreateFromFloat()
+    # Testing with float values
+    assert node.create_list(item_0=1.5, item_1=2.5, _dynamic_number=2) == ([1.5, 2.5],)
+
+    # Testing with float-convertible values
+    assert node.create_list(item_0=1, item_1="2.5", _dynamic_number=2) == ([1.0, 2.5],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_int():
+    node = DataListCreateFromInt()
+    # Testing with integer values
+    assert node.create_list(item_0=1, item_1=2, _dynamic_number=2) == ([1, 2],)
+
+    # Testing with int-convertible values
+    assert node.create_list(item_0="1", item_1=2.0, _dynamic_number=2) == ([1, 2],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_string():
+    node = DataListCreateFromString()
+    # Testing with string values
+    assert node.create_list(item_0="hello", item_1="world", _dynamic_number=2) == (["hello", "world"],)
+
+    # Testing with string-convertible values
+    assert node.create_list(item_0=123, item_1=True, _dynamic_number=2) == (["123", "True"],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create():
+    node = DataListCreate()
+    # Testing with one item
+    assert node.create_list(item_0="test", _dynamic_number=1) == (["test"],)
+
+    # Testing with multiple items of different types
+    assert node.create_list(item_0=1, item_1="two", item_2=3.0, _dynamic_number=3) == ([1, "two", 3.0],)
+
+    # Testing with empty list (no items)
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_boolean():
+    node = DataListCreateFromBoolean()
+    # Testing with boolean values
+    assert node.create_list(item_0=True, item_1=False, _dynamic_number=2) == ([True, False],)
+
+    # Testing with boolean-convertible values
+    assert node.create_list(item_0=1, item_1=0, _dynamic_number=2) == ([True, False],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_float():
+    node = DataListCreateFromFloat()
+    # Testing with float values
+    assert node.create_list(item_0=1.5, item_1=2.5, _dynamic_number=2) == ([1.5, 2.5],)
+
+    # Testing with float-convertible values
+    assert node.create_list(item_0=1, item_1="2.5", _dynamic_number=2) == ([1.0, 2.5],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_int():
+    node = DataListCreateFromInt()
+    # Testing with integer values
+    assert node.create_list(item_0=1, item_1=2, _dynamic_number=2) == ([1, 2],)
+
+    # Testing with int-convertible values
+    assert node.create_list(item_0="1", item_1=2.0, _dynamic_number=2) == ([1, 2],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
+def test_create_from_string():
+    node = DataListCreateFromString()
+    # Testing with string values
+    assert node.create_list(item_0="hello", item_1="world", _dynamic_number=2) == (["hello", "world"],)
+
+    # Testing with string-convertible values
+    assert node.create_list(item_0=123, item_1=True, _dynamic_number=2) == (["123", "True"],)
+
+    # Testing with empty list
+    assert node.create_list(_dynamic_number=0) == ([],)
+
+
 def test_pop_random():
    node = DataListPopRandom()
    result_list, item = node.pop_random_element(list=[1])
@@ -201,3 +382,75 @@ def test_last():
    assert node.get_last_element(list=["a", "b", "c"]) == ("c",)
    assert node.get_last_element(list=[]) == (None,)  # Empty list

+
+def test_to_list():
+    node = DataListToList()
+    # Test with regular list
+    assert node.convert(list=[1, 2, 3]) == ([1, 2, 3],)
+    # Test with empty list
+    assert node.convert(list=[]) == ([],)
+    # Test with mixed types
+    assert node.convert(list=[1, "two", 3.0]) == ([1, "two", 3.0],)
+
+
+def test_to_set():
+    node = DataListToSet()
+    # Test with regular list
+    assert node.convert(list=[1, 2, 3]) == ({1, 2, 3},)
+    # Test with empty list
+    assert node.convert(list=[]) == (set(),)
+    # Test with duplicates
+    assert node.convert(list=[1, 2, 1, 3, 2]) == ({1, 2, 3},)
+    # Test with mixed types (that can be in a set)
+    assert node.convert(list=[1, "two", 3.0]) == ({1, "two", 3.0},)
+
+
+def test_range():
+    node = DataListRange()
+    # Test with default start (0) and step (1)
+    assert node.create_range(stop=5) == ([0, 1, 2, 3, 4],)
+
+    # Test with custom start and stop
+    assert node.create_range(start=2, stop=6) == ([2, 3, 4, 5],)
+
+    # Test with custom step
+    assert node.create_range(start=0, stop=10, step=2) == ([0, 2, 4, 6, 8],)
+
+    # Test with negative numbers
+    assert node.create_range(start=-3, stop=3) == ([-3, -2, -1, 0, 1, 2],)
+
+    # Test backward counting
+    assert node.create_range(start=5, stop=0, step=-1) == ([5, 4, 3, 2, 1],)
+
+    # Test empty range
+    assert node.create_range(start=0, stop=0) == ([],)
+
+    # Test when start > stop with positive step (returns empty list)
+    assert node.create_range(start=10, stop=5) == ([],)
+
+    # Test with ValueError (step cannot be zero)
+    with pytest.raises(ValueError, match="Step cannot be zero"):
+        node.create_range(start=0, stop=5, step=0)
+
+
+def test_to_list():
+    node = DataListToList()
+    # Test with regular list
+    assert node.convert(list=[1, 2, 3]) == ([1, 2, 3],)
+    # Test with empty list
+    assert node.convert(list=[]) == ([],)
+    # Test with mixed types
+    assert node.convert(list=[1, "two", 3.0]) == ([1, "two", 3.0],)
+
+
+def test_to_set():
+    node = DataListToSet()
+    # Test with regular list
+    assert node.convert(list=[1, 2, 3]) == ({1, 2, 3},)
+    # Test with empty list
+    assert node.convert(list=[]) == (set(),)
+    # Test with duplicates
+    assert node.convert(list=[1, 2, 1, 3, 2]) == ({1, 2, 3},)
+    # Test with mixed types (that can be in a set)
+    assert node.convert(list=[1, "two", 3.0]) == ({1, "two", 3.0},)
+
@@ -19,6 +19,7 @@ from src.basic_data_handling.list_nodes import (
    ListMin,
    ListPop,
    ListPopRandom,
+    ListRange,
    ListRemove,
    ListReverse,
    ListSetItem,
@@ -221,3 +222,28 @@ def test_list_to_set():
    assert node.convert([1, 2, 3, 2, 1]) == ({1, 2, 3},)
    assert node.convert([]) == (set(),)
    assert node.convert(["a", "b", "a", "c"]) == ({"a", "b", "c"},)
+
+
+def test_list_range():
+    node = ListRange()
+    # Basic range with default step=1
+    assert node.create_range(start=0, stop=5) == ([0, 1, 2, 3, 4],)
+
+    # Range with custom step
+    assert node.create_range(start=0, stop=10, step=2) == ([0, 2, 4, 6, 8],)
+
+    # Negative step (counting down)
+    assert node.create_range(start=5, stop=0, step=-1) == ([5, 4, 3, 2, 1],)
+
+    # Start > stop with positive step (empty list)
+    assert node.create_range(start=10, stop=5) == ([],)
+
+    # Empty range
+    assert node.create_range(start=0, stop=0) == ([],)
+
+    # Test with negative indices
+    assert node.create_range(start=-5, stop=0) == ([-5, -4, -3, -2, -1],)
+
+    # Test error case: step=0
+    with pytest.raises(ValueError, match="Step cannot be zero"):
+        node.create_range(start=0, stop=5, step=0)