diff --git a/src/comfymath/control.py b/src/comfymath/control.py index 7927db0..c9bc704 100644 --- a/src/comfymath/control.py +++ b/src/comfymath/control.py @@ -1,3 +1,45 @@ from typing import Any, Mapping -NODE_CLASS_MAPPINGS: Mapping[str, Any] = {} + +class ChooseInt: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "condition": ("BOOLEAN", {"default": False}), + "a": ("INT", {"default": 0}), + "b": ("INT", {"default": 0}), + } + } + + RETURN_TYPES = ("INT",) + FUNCTION = "op" + CATEGORY = "math/control" + + def op(self, condition: bool, a: int, b: int) -> tuple[int]: + return (a if condition else b,) + + +class ChooseFloat: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "condition": ("BOOLEAN", {"default": False}), + "a": ("FLOAT", {"default": 0.0, "step": 0.001, "round": False}), + "b": ("FLOAT", {"default": 0.0, "step": 0.001, "round": False}), + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/control" + + def op(self, condition: bool, a: float, b: float) -> tuple[float]: + return (a if condition else b,) + + +NODE_CLASS_MAPPINGS: Mapping[str, Any] = { + "CM_ChooseInt": ChooseInt, + "CM_ChooseFloat": ChooseFloat, +} diff --git a/src/comfymath/float.py b/src/comfymath/float.py index 0783924..f7f3b02 100644 --- a/src/comfymath/float.py +++ b/src/comfymath/float.py @@ -151,9 +151,93 @@ def op(self, op: str, a: float, b: float) -> tuple[bool]: return (FLOAT_BINARY_CONDITIONS[op](a, b),) +def _float_fallback_binary( + *, fallback_mode: str, a: float, b: float, fallback_value: float +) -> float: + if fallback_mode == "A": + return a + if fallback_mode == "B": + return b + return fallback_value + + +def _float_fallback_unary(*, fallback_mode: str, a: float, fallback_value: float) -> float: + if fallback_mode == "A": + return a + return fallback_value + + +class FloatUnaryOperationConditional: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "condition": ("BOOLEAN", {"default": False}), + "fallback_mode": (["A", "constant"],), + "fallback_value": DEFAULT_FLOAT, + "op": (list(FLOAT_UNARY_OPERATIONS.keys()),), + "a": DEFAULT_FLOAT, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/float" + + def op( + self, + condition: bool, + fallback_mode: str, + fallback_value: float, + op: str, + a: float, + ) -> tuple[float]: + if condition: + return (FLOAT_UNARY_OPERATIONS[op](a),) + return (_float_fallback_unary(fallback_mode=fallback_mode, a=a, fallback_value=fallback_value),) + + +class FloatBinaryOperationConditional: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "condition": ("BOOLEAN", {"default": False}), + "fallback_mode": (["A", "B", "constant"],), + "fallback_value": DEFAULT_FLOAT, + "op": (list(FLOAT_BINARY_OPERATIONS.keys()),), + "a": DEFAULT_FLOAT, + "b": DEFAULT_FLOAT, + } + } + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/float" + + def op( + self, + condition: bool, + fallback_mode: str, + fallback_value: float, + op: str, + a: float, + b: float, + ) -> tuple[float]: + if condition: + return (FLOAT_BINARY_OPERATIONS[op](a, b),) + return ( + _float_fallback_binary( + fallback_mode=fallback_mode, a=a, b=b, fallback_value=fallback_value + ), + ) + + NODE_CLASS_MAPPINGS = { "CM_FloatUnaryOperation": FloatUnaryOperation, "CM_FloatUnaryCondition": FloatUnaryCondition, "CM_FloatBinaryOperation": FloatBinaryOperation, "CM_FloatBinaryCondition": FloatBinaryCondition, + "CM_FloatUnaryOperationConditional": FloatUnaryOperationConditional, + "CM_FloatBinaryOperationConditional": FloatBinaryOperationConditional, } diff --git a/src/comfymath/int.py b/src/comfymath/int.py index 576fc6a..d50ff55 100644 --- a/src/comfymath/int.py +++ b/src/comfymath/int.py @@ -75,7 +75,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: "required": {"op": (list(INT_UNARY_CONDITIONS.keys()),), "a": DEFAULT_INT} } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/int" @@ -113,7 +113,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/int" @@ -121,9 +121,91 @@ def op(self, op: str, a: int, b: int) -> tuple[bool]: return (INT_BINARY_CONDITIONS[op](a, b),) +def _int_fallback_binary(*, fallback_mode: str, a: int, b: int, fallback_value: int) -> int: + if fallback_mode == "A": + return a + if fallback_mode == "B": + return b + return fallback_value + + +def _int_fallback_unary(*, fallback_mode: str, a: int, fallback_value: int) -> int: + if fallback_mode == "A": + return a + return fallback_value + + +class IntUnaryOperationConditional: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "condition": ("BOOLEAN", {"default": False}), + "fallback_mode": (["A", "constant"],), + "fallback_value": DEFAULT_INT, + "op": (list(INT_UNARY_OPERATIONS.keys()),), + "a": DEFAULT_INT, + } + } + + RETURN_TYPES = ("INT",) + FUNCTION = "op" + CATEGORY = "math/int" + + def op( + self, + condition: bool, + fallback_mode: str, + fallback_value: int, + op: str, + a: int, + ) -> tuple[int]: + if condition: + return (INT_UNARY_OPERATIONS[op](a),) + return (_int_fallback_unary(fallback_mode=fallback_mode, a=a, fallback_value=fallback_value),) + + +class IntBinaryOperationConditional: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return { + "required": { + "condition": ("BOOLEAN", {"default": False}), + "fallback_mode": (["A", "B", "constant"],), + "fallback_value": DEFAULT_INT, + "op": (list(INT_BINARY_OPERATIONS.keys()),), + "a": DEFAULT_INT, + "b": DEFAULT_INT, + } + } + + RETURN_TYPES = ("INT",) + FUNCTION = "op" + CATEGORY = "math/int" + + def op( + self, + condition: bool, + fallback_mode: str, + fallback_value: int, + op: str, + a: int, + b: int, + ) -> tuple[int]: + if condition: + return (INT_BINARY_OPERATIONS[op](a, b),) + return ( + _int_fallback_binary( + fallback_mode=fallback_mode, a=a, b=b, fallback_value=fallback_value + ), + ) + + NODE_CLASS_MAPPINGS = { "CM_IntUnaryOperation": IntUnaryOperation, "CM_IntUnaryCondition": IntUnaryCondition, "CM_IntBinaryOperation": IntBinaryOperation, "CM_IntBinaryCondition": IntBinaryCondition, + "CM_IntUnaryOperationConditional": IntUnaryOperationConditional, + "CM_IntBinaryOperationConditional": IntBinaryOperationConditional, } diff --git a/src/comfymath/number.py b/src/comfymath/number.py index d2d3868..8dd36c6 100644 --- a/src/comfymath/number.py +++ b/src/comfymath/number.py @@ -40,7 +40,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/Number" @@ -78,7 +78,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/float" diff --git a/src/comfymath/vec.py b/src/comfymath/vec.py index 0b7180f..20e7e1e 100644 --- a/src/comfymath/vec.py +++ b/src/comfymath/vec.py @@ -125,7 +125,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/vec2" @@ -184,7 +184,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/vec2" @@ -257,7 +257,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/vec3" @@ -316,7 +316,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/vec3" @@ -389,7 +389,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/vec4" @@ -448,7 +448,7 @@ def INPUT_TYPES(cls) -> Mapping[str, Any]: } } - RETURN_TYPES = ("BOOL",) + RETURN_TYPES = ("BOOLEAN",) FUNCTION = "op" CATEGORY = "math/vec4"