From f8026c84a13619b0a58135539e6a334b1f0f0926 Mon Sep 17 00:00:00 2001 From: Evan Spearman Date: Wed, 19 Jul 2023 22:19:12 -0500 Subject: [PATCH] chore: Refactored int and float nodes, added number nodes --- __init__.py | 2 + src/comfymath/convert.py | 79 ++++-- src/comfymath/float.py | 588 +++++++++++---------------------------- src/comfymath/int.py | 356 ++++++++++-------------- src/comfymath/number.py | 205 ++++++++++++++ 5 files changed, 570 insertions(+), 660 deletions(-) create mode 100644 src/comfymath/number.py diff --git a/__init__.py b/__init__.py index ad2f934..264be62 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,7 @@ from .src.comfymath.convert import NODE_CLASS_MAPPINGS as convert_NCM from .src.comfymath.int import NODE_CLASS_MAPPINGS as int_NCM from .src.comfymath.float import NODE_CLASS_MAPPINGS as float_NCM +from .src.comfymath.number import NODE_CLASS_MAPPINGS as number_NCM from .src.comfymath.vec2 import NODE_CLASS_MAPPINGS as vec2_NCM from .src.comfymath.vec3 import NODE_CLASS_MAPPINGS as vec3_NCM from .src.comfymath.vec4 import NODE_CLASS_MAPPINGS as vec4_NCM @@ -11,6 +12,7 @@ NODE_CLASS_MAPPINGS = { **convert_NCM, **int_NCM, **float_NCM, + **number_NCM, **vec2_NCM, **vec3_NCM, **vec4_NCM, diff --git a/src/comfymath/convert.py b/src/comfymath/convert.py index 69319ca..b6ebbf4 100644 --- a/src/comfymath/convert.py +++ b/src/comfymath/convert.py @@ -3,12 +3,10 @@ from typing import Any, Mapping from .vec2 import Vec2, VEC2_ZERO from .vec3 import Vec3, VEC3_ZERO from .vec4 import Vec4, VEC4_ZERO +from .number import number class FloatToInt: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return {"required": {"a": ("FLOAT", {"default": 0.0})}} @@ -22,9 +20,6 @@ class FloatToInt: class IntToFloat: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return {"required": {"a": ("INT", {"default": 0})}} @@ -37,10 +32,57 @@ class IntToFloat: return (float(a),) -class ComposeVec2: - def __init__(self): - pass +class IntToNumber: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return {"required": {"a": ("INT", {"default": 0})}} + RETURN_TYPES = ("NUMBER",) + FUNCTION = "op" + CATEGORY = "math/conversion" + + def op(self, a: int) -> tuple[number]: + return (a,) + +class NumberToInt: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return {"required": {"a": ("NUMBER", {"default": 0.0})}} + + RETURN_TYPES = ("INT",) + FUNCTION = "op" + CATEGORY = "math/conversion" + + def op(self, a: number) -> tuple[int]: + return (int(a),) + + +class FloatToNumber: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return {"required": {"a": ("FLOAT", {"default": 0.0})}} + + RETURN_TYPES = ("NUMBER",) + FUNCTION = "op" + CATEGORY = "math/conversion" + + def op(self, a: float) -> tuple[number]: + return (a,) + +class NumberToFloat: + @classmethod + def INPUT_TYPES(cls) -> Mapping[str, Any]: + return {"required": {"a": ("NUMBER", {"default": 0.0})}} + + RETURN_TYPES = ("FLOAT",) + FUNCTION = "op" + CATEGORY = "math/conversion" + + def op(self, a: number) -> tuple[float]: + return (float(a),) + + +class ComposeVec2: @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return { @@ -59,9 +101,6 @@ class ComposeVec2: class BreakoutVec2: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return {"required": {"a": ("VEC2", {"default": VEC2_ZERO})}} @@ -75,9 +114,6 @@ class BreakoutVec2: class ComposeVec3: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return { @@ -97,9 +133,6 @@ class ComposeVec3: class BreakoutVec3: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return {"required": {"a": ("VEC3", {"default": VEC3_ZERO})}} @@ -113,9 +146,6 @@ class BreakoutVec3: class ComposeVec4: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return { @@ -136,9 +166,6 @@ class ComposeVec4: class BreakoutVec4: - def __init__(self): - pass - @classmethod def INPUT_TYPES(cls) -> Mapping[str, Any]: return {"required": {"a": ("VEC4", {"default": VEC4_ZERO})}} @@ -154,6 +181,10 @@ class BreakoutVec4: NODE_CLASS_MAPPINGS = { "FloatToInt": FloatToInt, "IntToFloat": IntToFloat, + "IntToNumber": IntToNumber, + "NumberToInt": NumberToInt, + "FloatToNumber": FloatToNumber, + "NumberToFloat": NumberToFloat, "ComposeVec2": ComposeVec2, "ComposeVec3": ComposeVec3, "ComposeVec4": ComposeVec4, diff --git a/src/comfymath/float.py b/src/comfymath/float.py index c1f27b9..5828aa0 100644 --- a/src/comfymath/float.py +++ b/src/comfymath/float.py @@ -1,440 +1,188 @@ import math -from abc import ABC, abstractmethod -from typing import Any, Mapping +from dataclasses import dataclass +from typing import Callable -class FloatUnaryOperator(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("FLOAT", {"default": 0.0})}} - - RETURN_TYPES = ("FLOAT",) - FUNCTION = "op" - - @abstractmethod - def op(self, a: float) -> tuple[float]: - pass +@dataclass +class FloatUnaryOperation: + name: str + function: Callable[[float], float] -class FloatBinaryOperator(ABC): - def __init__(self): - pass +@dataclass +class FloatUnaryCondition: + name: str + function: Callable[[float], bool] - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { + +@dataclass +class FloatBinaryOperation: + name: str + function: Callable[[float, float], float] + + +@dataclass +class FloatBinaryCondition: + name: str + function: Callable[[float, float], bool] + + +FLOAT_UNARY_OPERATIONS = [ + FloatUnaryOperation("Neg", lambda a: -a), + FloatUnaryOperation("Inc", lambda a: a + 1), + FloatUnaryOperation("Dec", lambda a: a - 1), + FloatUnaryOperation("Abs", lambda a: abs(a)), + FloatUnaryOperation("Sqr", lambda a: a * a), + FloatUnaryOperation("Cube", lambda a: a * a * a), + FloatUnaryOperation("Sqrt", lambda a: math.sqrt(a)), + FloatUnaryOperation("Exp", lambda a: math.exp(a)), + FloatUnaryOperation("Ln", lambda a: math.log(a)), + FloatUnaryOperation("Log10", lambda a: math.log10(a)), + FloatUnaryOperation("Log2", lambda a: math.log2(a)), + FloatUnaryOperation("Sin", lambda a: math.sin(a)), + FloatUnaryOperation("Cos", lambda a: math.cos(a)), + FloatUnaryOperation("Tan", lambda a: math.tan(a)), + FloatUnaryOperation("Asin", lambda a: math.asin(a)), + FloatUnaryOperation("Acos", lambda a: math.acos(a)), + FloatUnaryOperation("Atan", lambda a: math.atan(a)), + FloatUnaryOperation("Sinh", lambda a: math.sinh(a)), + FloatUnaryOperation("Cosh", lambda a: math.cosh(a)), + FloatUnaryOperation("Tanh", lambda a: math.tanh(a)), + FloatUnaryOperation("Asinh", lambda a: math.asinh(a)), + FloatUnaryOperation("Acosh", lambda a: math.acosh(a)), + FloatUnaryOperation("Atanh", lambda a: math.atanh(a)), + FloatUnaryOperation("Round", lambda a: round(a)), + FloatUnaryOperation("Floor", lambda a: math.floor(a)), + FloatUnaryOperation("Ceil", lambda a: math.ceil(a)), + FloatUnaryOperation("Trunc", lambda a: math.trunc(a)), + FloatUnaryOperation("Erf", lambda a: math.erf(a)), + FloatUnaryOperation("Erfc", lambda a: math.erfc(a)), + FloatUnaryOperation("Gamma", lambda a: math.gamma(a)), + FloatUnaryOperation("Radians", lambda a: math.radians(a)), + FloatUnaryOperation("Degrees", lambda a: math.degrees(a)), +] + +FLOAT_UNARY_CONDITIONS = [ + FloatUnaryCondition("IsZero", lambda a: a == 0.0), + FloatUnaryCondition("IsPositive", lambda a: a > 0.0), + FloatUnaryCondition("IsNegative", lambda a: a < 0.0), + FloatUnaryCondition("IsNonZero", lambda a: a != 0.0), + FloatUnaryCondition("IsPositiveInfinity", lambda a: math.isinf(a) and a > 0.0), + FloatUnaryCondition("IsNegativeInfinity", lambda a: math.isinf(a) and a < 0.0), + FloatUnaryCondition("IsNaN", lambda a: math.isnan(a)), + FloatUnaryCondition("IsFinite", lambda a: math.isfinite(a)), + FloatUnaryCondition("IsInfinite", lambda a: math.isinf(a)), + FloatUnaryCondition("IsEven", lambda a: a % 2 == 0.0), + FloatUnaryCondition("IsOdd", lambda a: a % 2 != 0.0), +] + +FLOAT_BINARY_OPERATIONS = [ + FloatBinaryOperation("Add", lambda a, b: a + b), + FloatBinaryOperation("Sub", lambda a, b: a - b), + FloatBinaryOperation("Mul", lambda a, b: a * b), + FloatBinaryOperation("Div", lambda a, b: a / b), + FloatBinaryOperation("Mod", lambda a, b: a % b), + FloatBinaryOperation("Pow", lambda a, b: a**b), + FloatBinaryOperation("FloorDiv", lambda a, b: a // b), + FloatBinaryOperation("Max", lambda a, b: max(a, b)), + FloatBinaryOperation("Min", lambda a, b: min(a, b)), + FloatBinaryOperation("Log", lambda a, b: math.log(a, b)), + FloatBinaryOperation("Atan2", lambda a, b: math.atan2(a, b)), +] + +FLOAT_BINARY_CONDITIONS = [ + FloatBinaryCondition("Eq", lambda a, b: a == b), + FloatBinaryCondition("Neq", lambda a, b: a != b), + FloatBinaryCondition("Gt", lambda a, b: a > b), + FloatBinaryCondition("Gte", lambda a, b: a >= b), + FloatBinaryCondition("Lt", lambda a, b: a < b), + FloatBinaryCondition("Lte", lambda a, b: a <= b), +] + + +def _get_float_unary_op_node_class(op: FloatUnaryOperation) -> type: + name = f"Float{op.name}" + class_dict = { + "INPUT_TYPES": lambda: {"required": {"a": ("FLOAT", {"default": 0.0})}}, + "RETURN_TYPES": ("FLOAT",), + "FUNCTION": "op", + "CATEGORY": "math/float", + "op": op.function, + } + return type(name, (), class_dict) + + +def _get_float_unary_cond_node_class(op: FloatUnaryCondition) -> type: + name = f"Float{op.name}" + class_dict = { + "INPUT_TYPES": lambda: {"required": {"a": ("FLOAT", {"default": 0.0})}}, + "RETURN_TYPES": ("INT",), + "FUNCTION": "op", + "CATEGORY": "math/float", + "op": lambda a: int(op.function(a)), + } + return type(name, (), class_dict) + + +def _get_float_binary_op_node_class(op: FloatBinaryOperation) -> type: + name = f"Float{op.name}" + class_dict = { + "INPUT_TYPES": lambda: { "required": { "a": ("FLOAT", {"default": 0.0}), "b": ("FLOAT", {"default": 0.0}), } - } - - RETURN_TYPES = ("FLOAT",) - FUNCTION = "op" - - @abstractmethod - def op(self, a: float, b: float) -> tuple[float]: - pass + }, + "RETURN_TYPES": ("FLOAT",), + "FUNCTION": "op", + "CATEGORY": "math/float", + "op": op.function, + } + return type(name, (), class_dict) -class FloatUnaryQuery(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("FLOAT", {"default": 0.0})}} - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - @abstractmethod - def op(self, a: float) -> tuple[int]: - pass - - -class FloatBinaryQuery(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { +def _get_float_binary_cond_node_class(op: FloatBinaryCondition) -> type: + name = f"Float{op.name}" + class_dict = { + "INPUT_TYPES": lambda: { "required": { "a": ("FLOAT", {"default": 0.0}), "b": ("FLOAT", {"default": 0.0}), } - } - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - @abstractmethod - def op(self, a: float, b: float) -> tuple[int]: - pass - - -class FloatAdd(FloatBinaryOperator): - def op(self, a: float, b: float) -> tuple[float]: - return (a + b,) - - CATEGORY = "math/float" - - -class FloatSub(FloatBinaryOperator): - def op(self, a: float, b: float) -> tuple[float]: - return (a - b,) - - CATEGORY = "math/float" - - -class FloatMul(FloatBinaryOperator): - def op(self, a: float, b: float) -> tuple[float]: - return (a * b,) - - CATEGORY = "math/float" - - -class FloatDiv(FloatBinaryOperator): - def op(self, a: float, b: float) -> tuple[float]: - return (a / b,) - - CATEGORY = "math/float" - - -class FloatLt(FloatBinaryQuery): - def op(self, a: float, b: float) -> tuple[int]: - return (int(a < b),) - - CATEGORY = "math/float/compare" - - -class FloatGt(FloatBinaryQuery): - def op(self, a: float, b: float) -> tuple[int]: - return (int(a > b),) - - CATEGORY = "math/float/compare" - - -class FloatLe(FloatBinaryQuery): - def op(self, a: float, b: float) -> tuple[int]: - return (int(a <= b),) - - CATEGORY = "math/float/compare" - - -class FloatGe(FloatBinaryQuery): - def op(self, a: float, b: float) -> tuple[int]: - return (int(a >= b),) - - CATEGORY = "math/float/compare" - - -class FloatEq(FloatBinaryQuery): - def op(self, a: float, b: float) -> tuple[int]: - return (int(a == b),) - - CATEGORY = "math/float/compare" - - -class FloatNe(FloatBinaryQuery): - def op(self, a: float, b: float) -> tuple[int]: - return (int(a != b),) - - CATEGORY = "math/float/compare" - - -class FloatNeg(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (-a,) - - CATEGORY = "math/float" - - -class FloatInc(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (a + 1,) - - CATEGORY = "math/float" - - -class FloatDec(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (a - 1,) - - CATEGORY = "math/float" - - -class FloatAbs(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.fabs(a),) - - CATEGORY = "math/float" - - -class FloatSqrt(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (a**0.5,) - - CATEGORY = "math/float" - - -class FloatPow(FloatBinaryOperator): - def op(self, a: float, b: float) -> tuple[float]: - return (a**b,) - - CATEGORY = "math/float" - - -class FloatLog(FloatBinaryOperator): - def op(self, a: float, b: float) -> tuple[float]: - return (math.log(a, b),) - - CATEGORY = "math/float" - - -class FloatCeil(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.ceil(a),) - - CATEGORY = "math/float" - - -class FloatFloor(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.floor(a),) - - CATEGORY = "math/float" - - -class FloatRound(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (round(a),) - - CATEGORY = "math/float" - - -class FloatTrunc(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.trunc(a),) - - CATEGORY = "math/float" - - -class FloatSin(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.sin(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatCos(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.cos(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatTan(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.tan(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatAsin(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.asin(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatAcos(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.acos(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatAtan(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.atan(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatAtan2(FloatBinaryOperator): - def op(self, a: float, b: float) -> tuple[float]: - return (math.atan2(a, b),) - - CATEGORY = "math/float/trigonometry" - - -class FloatLn(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.log(a),) - - CATEGORY = "math/float" - - -class FloatLog10(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.log10(a),) - - CATEGORY = "math/float" - - -class FloatLog2(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.log2(a),) - - CATEGORY = "math/float" - - -class FloatSinh(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.sinh(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatCosh(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.cosh(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatTanh(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.tanh(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatAsinh(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.asinh(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatAcosh(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.acosh(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatAtanh(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.atanh(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatExp(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.exp(a),) - - CATEGORY = "math/float" - - -class FloatExpm1(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.expm1(a),) - - CATEGORY = "math/float/functions" - - -class FloatErf(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.erf(a),) - - CATEGORY = "math/float/functions" - - -class FloatErfc(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.erfc(a),) - - CATEGORY = "math/float/functions" - - -class FloatGamma(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.gamma(a),) - - CATEGORY = "math/float/functions" - - -class FloatRadians(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.radians(a),) - - CATEGORY = "math/float/trigonometry" - - -class FloatDegrees(FloatUnaryOperator): - def op(self, a: float) -> tuple[float]: - return (math.degrees(a),) - - CATEGORY = "math/float/trigonometry" - + }, + "RETURN_TYPES": ("INT",), + "FUNCTION": "op", + "CATEGORY": "math/float", + "op": lambda a, b: int(op.function(a, b)), + } + return type(name, (), class_dict) + + +FLOAT_UNARY_OPERATION_CLASS_MAPPINGS = { + f"Float{op.name}": _get_float_unary_op_node_class(op) + for op in FLOAT_UNARY_OPERATIONS +} + +FLOAT_UNARY_CONDITION_CLASS_MAPPINGS = { + f"Float{op.name}": _get_float_unary_cond_node_class(op) + for op in FLOAT_UNARY_CONDITIONS +} + + +FLOAT_BINARY_OPERATION_CLASS_MAPPINGS = { + f"Float{op.name}": _get_float_binary_op_node_class(op) + for op in FLOAT_BINARY_OPERATIONS +} + +FLOAT_BINARY_CONDITION_CLASS_MAPPINGS = { + f"Float{op.name}": _get_float_binary_cond_node_class(op) + for op in FLOAT_BINARY_CONDITIONS +} NODE_CLASS_MAPPINGS = { - "FloatAdd": FloatAdd, - "FloatSub": FloatSub, - "FloatMul": FloatMul, - "FloatDiv": FloatDiv, - "FloatLt": FloatLt, - "FloatGt": FloatGt, - "FloatLe": FloatLe, - "FloatGe": FloatGe, - "FloatEq": FloatEq, - "FloatNe": FloatNe, - "FloatNeg": FloatNeg, - "FloatInc": FloatInc, - "FloatDec": FloatDec, - "FloatAbs": FloatAbs, - "FloatAbs": FloatAbs, - "FloatSqrt": FloatSqrt, - "FloatPow": FloatPow, - "FloatLog": FloatLog, - "FloatCeil": FloatCeil, - "FloatCeil": FloatCeil, - "FloatFloor": FloatFloor, - "FloatRound": FloatRound, - "FloatTrunc": FloatTrunc, - "FloatSin": FloatSin, - "FloatCos": FloatCos, - "FloatCos": FloatCos, - "FloatTan": FloatTan, - "FloatAsin": FloatAsin, - "FloatAcos": FloatAcos, - "FloatAtan": FloatAtan, - "FloatAtan2": FloatAtan2, - "FloatAtan2": FloatAtan2, - "FloatLn": FloatLn, - "FloatLog10": FloatLog10, - "FloatLog2": FloatLog2, - "FloatSinh": FloatSinh, - "FloatCosh": FloatCosh, - "FloatCosh": FloatCosh, - "FloatTanh": FloatTanh, - "FloatAsinh": FloatAsinh, - "FloatAcosh": FloatAcosh, - "FloatAtanh": FloatAtanh, - "FloatExp": FloatExp, - "FloatExpm1": FloatExpm1, - "FloatExpm1": FloatExpm1, - "FloatErf": FloatErf, - "FloatErfc": FloatErfc, - "FloatGamma": FloatGamma, - "FloatRadians": FloatRadians, - "FloatDegrees": FloatDegrees, + **FLOAT_UNARY_OPERATION_CLASS_MAPPINGS, + **FLOAT_UNARY_CONDITION_CLASS_MAPPINGS, + **FLOAT_BINARY_OPERATION_CLASS_MAPPINGS, + **FLOAT_BINARY_CONDITION_CLASS_MAPPINGS, } diff --git a/src/comfymath/int.py b/src/comfymath/int.py index 74b238b..ebe91a7 100644 --- a/src/comfymath/int.py +++ b/src/comfymath/int.py @@ -1,227 +1,151 @@ import math -from abc import ABC, abstractmethod -from typing import Any, Mapping +from dataclasses import dataclass +from typing import Callable, TypeAlias -class IntUnaryOperator(ABC): - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return {"required": {"a": ("INT", {"default": 0})}} - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - @abstractmethod - def op(self, a: int) -> tuple[int]: - pass +@dataclass +class IntUnaryOperation: + name: str + function: Callable[[int], int] -class IntBinaryOperator(ABC): - def __init__(self): - pass +@dataclass +class IntUnaryCondition: + name: str + function: Callable[[int], bool] - @classmethod - def INPUT_TYPES(cls) -> Mapping[str, Any]: - return { + +@dataclass +class IntBinaryOperation: + name: str + function: Callable[[int, int], int] + + +@dataclass +class IntBinaryCondition: + name: str + function: Callable[[int, int], bool] + + +INT_UNARY_OPERATIONS = [ + IntUnaryOperation("Abs", lambda a: abs(a)), + IntUnaryOperation("Neg", lambda a: -a), + IntUnaryOperation("Inc", lambda a: a + 1), + IntUnaryOperation("Dec", lambda a: a - 1), + IntUnaryOperation("Sqr", lambda a: a * a), + IntUnaryOperation("Cube", lambda a: a * a * a), + IntUnaryOperation("Not", lambda a: ~a), + IntUnaryOperation("Factorial", lambda a: math.factorial(a)), +] + +INT_UNARY_CONDITIONS = [ + IntUnaryCondition("IsZero", lambda a: a == 0), + IntUnaryCondition("IsNonZero", lambda a: a != 0), + IntUnaryCondition("IsPositive", lambda a: a > 0), + IntUnaryCondition("IsNegative", lambda a: a < 0), + IntUnaryCondition("IsEven", lambda a: a % 2 == 0), + IntUnaryCondition("IsOdd", lambda a: a % 2 == 1), +] + +INT_BINARY_OPERATIONS = [ + IntBinaryOperation("Add", lambda a, b: a + b), + IntBinaryOperation("Sub", lambda a, b: a - b), + IntBinaryOperation("Mul", lambda a, b: a * b), + IntBinaryOperation("Div", lambda a, b: a // b), + IntBinaryOperation("Mod", lambda a, b: a % b), + IntBinaryOperation("Pow", lambda a, b: a**b), + IntBinaryOperation("And", lambda a, b: a & b), + IntBinaryOperation("Nand", lambda a, b: ~a & b), + IntBinaryOperation("Or", lambda a, b: a | b), + IntBinaryOperation("Nor", lambda a, b: ~a & b), + IntBinaryOperation("Xor", lambda a, b: a ^ b), + IntBinaryOperation("Xnor", lambda a, b: ~a ^ b), + IntBinaryOperation("Shl", lambda a, b: a << b), + IntBinaryOperation("Shr", lambda a, b: a >> b), + IntBinaryOperation("Max", lambda a, b: max(a, b)), + IntBinaryOperation("Min", lambda a, b: min(a, b)), +] + +INT_BINARY_CONDITIONS = [ + IntBinaryCondition("Eq", lambda a, b: a == b), + IntBinaryCondition("Neq", lambda a, b: a != b), + IntBinaryCondition("Gt", lambda a, b: a > b), + IntBinaryCondition("Lt", lambda a, b: a < b), + IntBinaryCondition("Geq", lambda a, b: a >= b), + IntBinaryCondition("Leq", lambda a, b: a <= b), +] + + +def _get_int_unary_op_node_class(op: IntUnaryOperation) -> type: + name = f"Int{op.name}" + class_dict = { + "INPUT_TYPES": lambda: {"required": {"a": ("INT", {"default": 0})}}, + "RETURN_TYPES": ("INT",), + "FUNCTION": "op", + "CATEGORY": "math/int", + "op": lambda a: op.function(a), + } + return type(name, (), class_dict) + +def _get_int_unary_cond_node_class(op: IntUnaryCondition) -> type: + name = f"Int{op.name}" + class_dict = { + "INPUT_TYPES": lambda: {"required": {"a": ("INT", {"default": 0})}}, + "RETURN_TYPES": ("INT",), + "FUNCTION": "op", + "CATEGORY": "math/int", + "op": lambda a: int(op.function(a)), + } + return type(name, (), class_dict) + + +def _get_int_binary_op_node_class(op: IntBinaryOperation) -> type: + name = f"Int{op.name}" + class_dict = { + "INPUT_TYPES": lambda: { "required": {"a": ("INT", {"default": 0}), "b": ("INT", {"default": 0})} - } - - RETURN_TYPES = ("INT",) - FUNCTION = "op" - - @abstractmethod - def op(self, a: int, b: int) -> tuple[int]: - pass - - -class IntAdd(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a + b,) - - CATEGORY = "math/int" - - -class IntSub(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a - b,) - - CATEGORY = "math/int" - - -class IntMul(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a * b,) - - CATEGORY = "math/int" - - -class IntDiv(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a // b,) - - CATEGORY = "math/int" - - -class IntMod(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a % b,) - - CATEGORY = "math/int" - - -class IntPow(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a**b,) - - CATEGORY = "math/int" - - -class IntLt(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (int(a < b),) - - CATEGORY = "math/int/compare" - - -class IntGt(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (int(a > b),) - - CATEGORY = "math/int/compare" - - -class IntLe(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (int(a <= b),) - - CATEGORY = "math/int/compare" - - -class IntGe(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (int(a >= b),) - - CATEGORY = "math/int/compare" - - -class IntEq(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (int(a == b),) - - CATEGORY = "math/int/compare" - - -class IntNe(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (int(a != b),) - - CATEGORY = "math/int/compare" - - -class IntAnd(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a & b,) - - CATEGORY = "math/int/bitwise" - - -class IntOr(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a | b,) - - CATEGORY = "math/int/bitwise" - - -class IntXor(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (a ^ b,) - - CATEGORY = "math/int/bitwise" - - -class IntXnor(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (~(a ^ b),) - - CATEGORY = "math/int/bitwise" - - -class IntNand(IntBinaryOperator): - def op(self, a: int, b: int) -> tuple[int]: - return (~(a & b),) - - CATEGORY = "math/int/bitwise" - - -class IntNot(IntUnaryOperator): - def op(self, a: int) -> tuple[int]: - return (~a,) - - CATEGORY = "math/int/bitwise" - - -class IntNeg(IntUnaryOperator): - def op(self, a: int) -> tuple[int]: - return (-a,) - - CATEGORY = "math/int" - - -class IntInc(IntUnaryOperator): - def op(self, a: int) -> tuple[int]: - return (a + 1,) - - CATEGORY = "math/int" - - -class IntDec(IntUnaryOperator): - def op(self, a: int) -> tuple[int]: - return (a - 1,) - - CATEGORY = "math/int" - - -class IntAbs(IntUnaryOperator): - def op(self, a: int) -> tuple[int]: - return (abs(a),) - - CATEGORY = "math/int" - - -class IntFactorial(IntUnaryOperator): - def op(self, a: int) -> tuple[int]: - return (math.factorial(a),) - - CATEGORY = "math/int" - + }, + "RETURN_TYPES": ("INT",), + "FUNCTION": "op", + "CATEGORY": "math/int", + "op": lambda a, b: op.function(a, b), + } + return type(name, (), class_dict) + +def _get_int_binary_cond_node_class(op: IntBinaryCondition) -> type: + name = f"Int{op.name}" + class_dict = { + "INPUT_TYPES": lambda: { + "required": {"a": ("INT", {"default": 0}), "b": ("INT", {"default": 0})} + }, + "RETURN_TYPES": ("INT",), + "FUNCTION": "op", + "CATEGORY": "math/int", + "op": lambda a, b: int(op.function(a, b)), + } + return type(name, (), class_dict) + + +FLOAT_UNARY_OPERATION_CLASS_MAPPINGS = { + f"Int{op.name}": _get_int_unary_op_node_class(op) for op in INT_UNARY_OPERATIONS +} + +FLOAT_UNARY_CONDITION_CLASS_MAPPINGS = { + f"Int{op.name}": _get_int_unary_cond_node_class(op) for op in INT_UNARY_CONDITIONS +} + +FLOAT_BINARY_OPERATION_CLASS_MAPPINGS = { + f"Int{op.name}": _get_int_binary_op_node_class(op) for op in INT_BINARY_OPERATIONS +} + +FLOAT_BINARY_CONDITION_CLASS_MAPPINGS = { + f"Int{op.name}": _get_int_binary_cond_node_class(op) for op in INT_BINARY_CONDITIONS +} NODE_CLASS_MAPPINGS = { - "IntAdd": IntAdd, - "IntSub": IntSub, - "IntMul": IntMul, - "IntDiv": IntDiv, - "IntMod": IntMod, - "IntPow": IntPow, - "IntLt": IntLt, - "IntGt": IntGt, - "IntLe": IntLe, - "IntGe": IntGe, - "IntEq": IntEq, - "IntNe": IntNe, - "IntAnd": IntAnd, - "IntOr": IntOr, - "IntXor": IntXor, - "IntXnor": IntXnor, - "IntNand": IntNand, - "IntNot": IntNot, - "IntNeg": IntNeg, - "IntInc": IntInc, - "IntInc": IntInc, - "IntDec": IntDec, - "IntAbs": IntAbs, - "IntFactorial": IntFactorial, + **FLOAT_UNARY_OPERATION_CLASS_MAPPINGS, + **FLOAT_UNARY_CONDITION_CLASS_MAPPINGS, + **FLOAT_BINARY_OPERATION_CLASS_MAPPINGS, + **FLOAT_BINARY_CONDITION_CLASS_MAPPINGS, } diff --git a/src/comfymath/number.py b/src/comfymath/number.py new file mode 100644 index 0000000..cfed571 --- /dev/null +++ b/src/comfymath/number.py @@ -0,0 +1,205 @@ +from dataclasses import dataclass +from typing import Callable, TypeAlias, Sequence + +from .int import ( + INT_UNARY_OPERATIONS, + INT_UNARY_CONDITIONS, + INT_BINARY_OPERATIONS, + INT_BINARY_CONDITIONS, +) +from .float import ( + FLOAT_UNARY_OPERATIONS, + FLOAT_UNARY_CONDITIONS, + FLOAT_BINARY_OPERATIONS, + FLOAT_BINARY_CONDITIONS, + FloatUnaryOperation, + FloatUnaryCondition, + FloatBinaryOperation, + FloatBinaryCondition, +) + +number: TypeAlias = int | float + + +@dataclass +class NumberUnaryOperation: + name: str + function: Callable[[number], number] + + +@dataclass +class NumberUnaryCondition: + name: str + function: Callable[[number], bool] + + +@dataclass +class NumberBinaryOperation: + name: str + function: Callable[[number, number], number] + + +@dataclass +class NumberBinaryCondition: + name: str + function: Callable[[number, number], bool] + + +def _float_unary_operation_to_num_unary_operation( + op: FloatUnaryOperation, +) -> NumberUnaryOperation: + return NumberUnaryOperation(op.name, op.function) + + +def _float_unary_condition_to_num_unary_condition( + op: FloatUnaryCondition, +) -> NumberUnaryCondition: + return NumberUnaryCondition(op.name, op.function) + + +def _float_binary_operation_to_num_binary_operation( + op: FloatBinaryOperation, +) -> NumberBinaryOperation: + return NumberBinaryOperation(op.name, op.function) + + +def _float_binary_condition_to_num_binary_condition( + op: FloatBinaryCondition, +) -> NumberBinaryCondition: + return NumberBinaryCondition(op.name, op.function) + + +def _combine_unary_operations() -> Sequence[NumberUnaryOperation]: + float_unary_op_names = {op.name for op in FLOAT_UNARY_OPERATIONS} + int_unary_op_names = {op.name for op in INT_UNARY_OPERATIONS} + num_unary_op_names = float_unary_op_names & int_unary_op_names + return [ + _float_unary_operation_to_num_unary_operation(op) + for op in FLOAT_UNARY_OPERATIONS + if op.name in num_unary_op_names + ] + + +def _combine_unary_conditions() -> Sequence[NumberUnaryCondition]: + float_unary_cond_names = {op.name for op in FLOAT_UNARY_CONDITIONS} + int_unary_cond_names = {op.name for op in INT_UNARY_CONDITIONS} + num_unary_cond_names = float_unary_cond_names & int_unary_cond_names + return [ + _float_unary_condition_to_num_unary_condition(op) + for op in FLOAT_UNARY_CONDITIONS + if op.name in num_unary_cond_names + ] + + +def _combine_binary_operations() -> Sequence[NumberBinaryOperation]: + float_binary_op_names = {op.name for op in FLOAT_BINARY_OPERATIONS} + int_binary_op_names = {op.name for op in INT_BINARY_OPERATIONS} + num_binary_op_names = float_binary_op_names & int_binary_op_names + return [ + _float_binary_operation_to_num_binary_operation(op) + for op in FLOAT_BINARY_OPERATIONS + if op.name in num_binary_op_names + ] + + +def _combine_binary_conditions() -> Sequence[NumberBinaryCondition]: + float_binary_cond_names = {op.name for op in FLOAT_BINARY_CONDITIONS} + int_binary_cond_names = {op.name for op in INT_BINARY_CONDITIONS} + num_binary_cond_names = float_binary_cond_names & int_binary_cond_names + return [ + _float_binary_condition_to_num_binary_condition(op) + for op in FLOAT_BINARY_CONDITIONS + if op.name in num_binary_cond_names + ] + + +def _get_number_unary_op_node_class(op: NumberUnaryOperation) -> type: + name = f"Number{op.name}" + class_dict = { + "INPUT_TYPES": lambda: {"required": {"a": ("NUMBER", {"default": 0.0})}}, + "RETURN_TYPES": ("NUMBER",), + "FUNCTION": "op", + "CATEGORY": "math/number", + "op": op.function, + } + return type(name, (), class_dict) + + +def _get_number_unary_cond_node_class(op: NumberUnaryCondition) -> type: + name = f"Number{op.name}" + class_dict = { + "INPUT_TYPES": lambda: {"required": {"a": ("NUMBER", {"default": 0.0})}}, + "RETURN_TYPES": ("INT",), + "FUNCTION": "op", + "CATEGORY": "math/number", + "op": lambda a: int(op.function(a)), + } + return type(name, (), class_dict) + + +def _get_number_binary_op_node_class(op: NumberBinaryOperation) -> type: + name = f"Number{op.name}" + class_dict = { + "INPUT_TYPES": lambda: { + "required": { + "a": ("NUMBER", {"default": 0.0}), + "b": ("NUMBER", {"default": 0.0}), + } + }, + "RETURN_TYPES": ("NUMBER",), + "FUNCTION": "op", + "CATEGORY": "math/number", + "op": op.function, + } + return type(name, (), class_dict) + + +def _get_number_binary_cond_node_class(op: NumberBinaryCondition) -> type: + name = f"Number{op.name}" + class_dict = { + "INPUT_TYPES": lambda: { + "required": { + "a": ("NUMBER", {"default": 0.0}), + "b": ("NUMBER", {"default": 0.0}), + } + }, + "RETURN_TYPES": ("NUMBER",), + "FUNCTION": "op", + "CATEGORY": "math/number", + "op": lambda a, b: int(op.function(a, b)), + } + return type(name, (), class_dict) + + +NUMBER_UNARY_OPERATIONS = _combine_unary_operations() +NUMBER_UNARY_CONDITIONS = _combine_unary_conditions() +NUMBER_BINARY_OPERATIONS = _combine_binary_operations() +NUMBER_BINARY_CONDITIONS = _combine_binary_conditions() + +NUMBER_UNARY_OPERATION_CLASS_MAPPINGS = { + f"Number{op.name}": _get_number_unary_op_node_class(op) + for op in NUMBER_UNARY_OPERATIONS +} + +NUMBER_UNARY_CONDITION_CLASS_MAPPINGS = { + f"Number{op.name}": _get_number_unary_cond_node_class(op) + for op in NUMBER_UNARY_CONDITIONS +} + + +NUMBER_BINARY_OPERATION_CLASS_MAPPINGS = { + f"Number{op.name}": _get_number_binary_op_node_class(op) + for op in NUMBER_BINARY_OPERATIONS +} + +NUMBER_BINARY_CONDITION_CLASS_MAPPINGS = { + f"Number{op.name}": _get_number_binary_cond_node_class(op) + for op in NUMBER_BINARY_CONDITIONS +} + +NODE_CLASS_MAPPINGS = { + **NUMBER_UNARY_OPERATION_CLASS_MAPPINGS, + **NUMBER_UNARY_CONDITION_CLASS_MAPPINGS, + **NUMBER_BINARY_OPERATION_CLASS_MAPPINGS, + **NUMBER_BINARY_CONDITION_CLASS_MAPPINGS, +}