Add flow selection nodes

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.

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",
 }

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",

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