diff --git a/src/basic_data_handling/control_flow_nodes.py b/src/basic_data_handling/control_flow_nodes.py index 92cd493..7d96b91 100644 --- a/src/basic_data_handling/control_flow_nodes.py +++ b/src/basic_data_handling/control_flow_nodes.py @@ -34,27 +34,8 @@ class IfElse(ComfyNodeABC): needed.append("if_false") return needed - def execute(self, **kwargs) -> tuple[Any]: - # Check if the main condition is true - if kwargs.get("if", False): - return (kwargs.get("then"),) - - # Check each elif condition - elif_index = 0 - while True: - elif_key = f"elif_{elif_index}" - then_key = f"then_{elif_index}" - - if elif_key not in kwargs: - break - - if kwargs.get(elif_key, False): - return (kwargs.get(then_key),) - - elif_index += 1 - - # If no conditions were true, return the else value - return (kwargs.get("else"),) + def execute(self, condition: bool, if_true: Any, if_false: Any) -> tuple[Any]: + return (if_true if condition else if_false,) class IfElifElse(ComfyNodeABC): diff --git a/tests/test_control_flow_nodes.py b/tests/test_control_flow_nodes.py index f3de291..de6072f 100644 --- a/tests/test_control_flow_nodes.py +++ b/tests/test_control_flow_nodes.py @@ -21,19 +21,19 @@ def test_switch_case(): node = SwitchCase() # Selector within range, valid cases - assert node.execute(0, "Case 0", "Case 1", "Case 2", "Case 3") == ("Case 0",) - assert node.execute(1, "Case 0", "Case 1", "Case 2", "Case 3") == ("Case 1",) - assert node.execute(3, "Case 0", "Case 1", "Case 2", "Case 3") == ("Case 3",) + assert node.execute(0, case_0="Case 0", case_1="Case 1", case_2="Case 2", case_3="Case 3") == ("Case 0",) + assert node.execute(1, case_0="Case 0", case_1="Case 1", case_2="Case 2", case_3="Case 3") == ("Case 1",) + assert node.execute(3, case_0="Case 0", case_1="Case 1", case_2="Case 2", case_3="Case 3") == ("Case 3",) # Selector out of range, fallback to default case - assert node.execute(4, "Case 0", "Case 1", "Case 2", "Case 3", default="Default") == ("Default",) - assert node.execute(99, "Case 0", "Case 1", "Case 2", "Case 3", default="Default") == ("Default",) + assert node.execute(4, case_0="Case 0", case_1="Case 1", case_2="Case 2", case_3="Case 3", default="Default") == ("Default",) + assert node.execute(99, case_0="Case 0", case_1="Case 1", case_2="Case 2", case_3="Case 3", default="Default") == ("Default",) # Selector out of range, no default provided - assert node.execute(4, "Case 0", "Case 1", "Case 2", "Case 3") == (None,) + assert node.execute(4, case_0="Case 0", case_1="Case 1", case_2="Case 2", case_3="Case 3") == (None,) # Edge cases - assert node.execute(0, None, "Case 1", "Case 2") == (None,) # Case with None - assert node.execute(2, "Case 0", "Case 1", "Case 2") == ("Case 2",) - assert node.execute(0, 123, False, "Case 2") == (123,) # Different data types - assert node.execute(10, "Case 0", "Case 1", "Case 2", default="Fallback") == ("Fallback",) + assert node.execute(0, case_0=None, case_1="Case 1", case_2="Case 2") == (None,) # Case with None + assert node.execute(2, case_0="Case 0", case_1="Case 1", case_2="Case 2") == ("Case 2",) + assert node.execute(0, case_0=123, case_1=False, case_2="Case 2") == (123,) # Different data types + assert node.execute(10, case_0="Case 0", case_1="Case 1", case_2="Case 2", default="Fallback") == ("Fallback",)