From 622f471bfac953fae15ac3fa07b1a18f4f6c0843 Mon Sep 17 00:00:00 2001 From: StableLlama Date: Mon, 9 Jun 2025 13:33:10 +0200 Subject: [PATCH] Add rage nodes, disable flow and force calculation nodes --- src/basic_data_handling/control_flow_nodes.py | 85 +++++- src/basic_data_handling/data_list_nodes.py | 33 +++ src/basic_data_handling/list_nodes.py | 33 ++- tests/test_data_list_nodes.py | 267 +++++++++++++++++- tests/test_list_nodes.py | 26 ++ 5 files changed, 427 insertions(+), 17 deletions(-) diff --git a/src/basic_data_handling/control_flow_nodes.py b/src/basic_data_handling/control_flow_nodes.py index a18d53a..772d622 100644 --- a/src/basic_data_handling/control_flow_nodes.py +++ b/src/basic_data_handling/control_flow_nodes.py @@ -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", } diff --git a/src/basic_data_handling/data_list_nodes.py b/src/basic_data_handling/data_list_nodes.py index 5da7b9c..1ecc315 100644 --- a/src/basic_data_handling/data_list_nodes.py +++ b/src/basic_data_handling/data_list_nodes.py @@ -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", diff --git a/src/basic_data_handling/list_nodes.py b/src/basic_data_handling/list_nodes.py index 5d3e3e8..a2be8f6 100644 --- a/src/basic_data_handling/list_nodes.py +++ b/src/basic_data_handling/list_nodes.py @@ -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", diff --git a/tests/test_data_list_nodes.py b/tests/test_data_list_nodes.py index ac0c5a5..797a8b9 100644 --- a/tests/test_data_list_nodes.py +++ b/tests/test_data_list_nodes.py @@ -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},) + diff --git a/tests/test_list_nodes.py b/tests/test_list_nodes.py index 668fba8..3bd732a 100644 --- a/tests/test_list_nodes.py +++ b/tests/test_list_nodes.py @@ -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)