diff --git a/pyproject.toml b/pyproject.toml index aae12b3..e669584 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "basic_data_handling" -version = "0.3.0" +version = "0.3.2" description = """NOTE: Still in development! Expect breaking changes! Basic Python functions for manipulating data that every programmer is used to. Currently supported ComfyUI data types: BOOLEAN, FLOAT, INT, STRING and data lists. diff --git a/src/basic_data_handling/control_flow_nodes.py b/src/basic_data_handling/control_flow_nodes.py index 7d96b91..88f343d 100644 --- a/src/basic_data_handling/control_flow_nodes.py +++ b/src/basic_data_handling/control_flow_nodes.py @@ -1,6 +1,8 @@ from typing import Any from inspect import cleandoc from comfy.comfy_types.node_typing import IO, ComfyNodeABC +from comfy_execution.graph import ExecutionBlocker + class IfElse(ComfyNodeABC): """ @@ -53,11 +55,11 @@ class IfElifElse(ComfyNodeABC): def INPUT_TYPES(cls): return { "required": { - "if": (IO.BOOLEAN, {}), + "if": (IO.BOOLEAN, {"forceInput": True}), "then": (IO.ANY, {"lazy": True}), }, "optional": { - "elif_0": (IO.BOOLEAN, {"lazy": True, "_dynamic": "number", "_dynamicGroup": 0}), + "elif_0": (IO.BOOLEAN, {"forceInput": True, "lazy": True, "_dynamic": "number", "_dynamicGroup": 0}), "then_0": (IO.ANY, {"lazy": True, "_dynamic": "number", "_dynamicGroup": 0}), "else": (IO.ANY, {"lazy": True}), } @@ -192,14 +194,76 @@ def execute(self, selector: int, **kwargs) -> tuple[Any]: return (kwargs.get("default"),) +class FlowSelect(ComfyNodeABC): + """ + Select the direction of the flow. + + This node takes a value and directs it to either the "true" or "false" output. + + Note: for dynamic switching in a Data Flow you might want to use + "filter select" instead. + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "value": (IO.ANY, {}), + "select": (IO.BOOLEAN, {}), + } + } + + RETURN_TYPES = (IO.ANY, IO.ANY) + RETURN_NAMES = ("true", "false") + CATEGORY = "Basic/flow control" + DESCRIPTION = cleandoc(__doc__ or "") + FUNCTION = "select" + + def select(self, value, select = True) -> tuple[Any, Any]: + if select: + return value, ExecutionBlocker(None) + else: + return ExecutionBlocker(None), value + + +class ExecutionOrder(ComfyNodeABC): + """ + Force execution order in the workflow. + + This node is lightweight and does not affect the workflow. It is used to force + the execution order of nodes in the workflow. You only need to chain this node + with the other execution order nodes in the desired order and add any + output of the nodes you want to force execution order on. + """ + @classmethod + def INPUT_TYPES(cls): + return { + "optional": { + "E/O": ("E/O", {}), + "any node output": (IO.ANY, {}), + } + } + + RETURN_TYPES = ("E/O",) + CATEGORY = "Basic/flow control" + DESCRIPTION = cleandoc(__doc__ or "") + FUNCTION = "execute" + + def execute(self, **kwargs) -> tuple[Any]: + return (None,) + + NODE_CLASS_MAPPINGS = { "Basic data handling: IfElse": IfElse, "Basic data handling: IfElifElse": IfElifElse, "Basic data handling: SwitchCase": SwitchCase, + "Basic data handling: FlowSelect": FlowSelect, + "Basic data handling: ExecutionOrder": ExecutionOrder, } 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: FlowSelect": "flow select", + "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 346ba61..003461e 100644 --- a/src/basic_data_handling/data_list_nodes.py +++ b/src/basic_data_handling/data_list_nodes.py @@ -295,6 +295,46 @@ def filter_data(self, **kwargs: list[Any]) -> tuple[list[Any]]: return (result,) +class DataListFilterSelect(ComfyNodeABC): + """ + Filters a Data List using boolean values. + + This node takes a value Data List and a filter Data List (containing only boolean values). + It returns two new Data Lists containing only the elements from the value list where the + corresponding element in the filter list is true or false. + + If the lists have different lengths, the last element of the shorter list is repeated + till the lengths are matching. + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "value": (IO.ANY, {}), + "select": (IO.BOOLEAN, {}), + } + } + + RETURN_TYPES = (IO.ANY, IO.ANY) + RETURN_NAMES = ("true", "false") + CATEGORY = "Basic/Data List" + DESCRIPTION = cleandoc(__doc__ or "") + FUNCTION = "select" + INPUT_IS_LIST = True + OUTPUT_IS_LIST = (True, True,) + + def select(self, **kwargs: list[Any]) -> tuple[list[Any]]: + values = kwargs.get('value', []) + selects = kwargs.get('select', []) + + # Create a new list with only items where the filter is False + result_true, result_false = [], [] + for value, select in zip(values, selects): + (result_true if select else result_false).append(value) + + return result_true, result_false + + class DataListGetItem(ComfyNodeABC): """ Retrieves an item at a specified position in a list. @@ -771,7 +811,7 @@ def INPUT_TYPES(cls): INPUT_IS_LIST = True def convert(self, **kwargs: list[Any]) -> tuple[list[Any]]: - return (kwargs.get('list', []).copy(),) + return (list(kwargs.get('list', [])).copy(),) class DataListToSet(ComfyNodeABC): @@ -810,6 +850,7 @@ def convert(self, **kwargs: list[Any]) -> tuple[set[Any]]: "Basic data handling: DataListCount": DataListCount, "Basic data handling: DataListExtend": DataListExtend, "Basic data handling: DataListFilter": DataListFilter, + "Basic data handling: DataListFilterSelect": DataListFilterSelect, "Basic data handling: DataListGetItem": DataListGetItem, "Basic data handling: DataListIndex": DataListIndex, "Basic data handling: DataListInsert": DataListInsert, @@ -838,6 +879,7 @@ def convert(self, **kwargs: list[Any]) -> tuple[set[Any]]: "Basic data handling: DataListCount": "count", "Basic data handling: DataListExtend": "extend", "Basic data handling: DataListFilter": "filter", + "Basic data handling: DataListFilterSelect": "filter select", "Basic data handling: DataListGetItem": "get item", "Basic data handling: DataListIndex": "index", "Basic data handling: DataListInsert": "insert", diff --git a/tests/conftest.py b/tests/conftest.py index aa09b9b..e4d9d2c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -26,3 +26,11 @@ sys.modules["comfy"] = mock_comfy sys.modules["comfy.comfy_types"] = mock_comfy sys.modules["comfy.comfy_types.node_typing"] = mock_comfy + +def mock_execution_blocker(_): + return None + +mock_comfy_execution = MagicMock() +mock_comfy_execution.ExecutionBlocker = mock_execution_blocker +sys.modules["comfy_execution"] = mock_comfy_execution +sys.modules["comfy_execution.graph"] = mock_comfy_execution