diff --git a/comfy_extras/nodes_math.py b/comfy_extras/nodes_math.py new file mode 100644 index 000000000..7d20a57bd --- /dev/null +++ b/comfy_extras/nodes_math.py @@ -0,0 +1,73 @@ +"""Math expression node using JSONata for evaluation. + +Provides a ComfyMathExpression node that evaluates JSONata expressions +against dynamically-grown numeric inputs. Supports frontend eager +evaluation via the eager_eval schema field. +""" + +from __future__ import annotations + +import jsonata as jsonata_lib +from typing_extensions import override + +from comfy_api.latest import ComfyExtension, io + + +class MathExpressionNode(io.ComfyNode): + """Evaluates a JSONata expression against dynamically-grown inputs.""" + + @classmethod + def define_schema(cls) -> io.Schema: + template = io.MatchType.Template( + "num", allowed_types=[io.Float, io.Int] + ) + autogrow = io.Autogrow.TemplatePrefix( + input=io.MatchType.Input("value", template=template), + prefix="value", + min=1, + max=100, + ) + return io.Schema( + node_id="ComfyMathExpression", + display_name="Math Expression", + category="math", + search_aliases=[ + "expression", "formula", "calculate", "eval", "math" + ], + inputs=[ + io.String.Input("expression", default="a + b"), + io.Autogrow.Input("values", template=autogrow), + ], + outputs=[ + io.MatchType.Output( + template=template, display_name="result" + ), + ], + is_output_node=True, + eager_eval=io.EagerEval(expr_widget="expression"), + ) + + @classmethod + def execute( + cls, expression: str, values: io.Autogrow.Type + ) -> io.NodeOutput: + context: dict = {} + for i, (key, val) in enumerate(values.items()): + context[chr(ord("a") + i)] = val # positional: a, b, c, ... + context[key] = val # also by input name: value0, value1, ... + context["values"] = list(values.values()) # for $sum(values) etc. + + result = jsonata_lib.Jsonata(expression).evaluate(context) + return io.NodeOutput( + result, ui={"result": [result], "context": [context]} + ) + + +class MathExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [MathExpressionNode] + + +async def comfy_entrypoint() -> MathExtension: + return MathExtension()