Cache parsed expressions

check_lazy_status and execute parsed the same expression twice per run.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

more_math/ScriptTextWindow.py

@@ -22,4 +22,4 @@ class ScriptTextInput(io.ComfyNode):
 
     @classmethod
     def execute(cls, script):
-        return (remove_dead_code(parse_expr(script)),)
+        return (remove_dead_code(parse_expr.__wrapped__(script)),)

more_math/helper_functions.py

@@ -1,3 +1,5 @@
+import functools
+
 from antlr4.error.ErrorListener import ErrorListener
 from antlr4 import InputStream, CommonTokenStream
 from antlr4.atn.PredictionMode import PredictionMode
@@ -66,8 +68,10 @@ def as_tensor(value, shape):
     return torch.cat(value)
 
 
+# check_lazy_status and execute both parse the same expression; trees are only read, so they can be shared.
+@functools.lru_cache(maxsize=64)
 def parse_expr(expr: str):
-    """Parse a math expression and return the parse tree."""
+    """Parse a math expression and return the parse tree. The tree is shared, use parse_expr.__wrapped__ for one you modify."""
     input_stream = InputStream(expr)
     lexer = MathExprLexer(input_stream)
     stream = CommonTokenStream(lexer)

tests/test_dead_code.py

@@ -9,7 +9,7 @@ from more_math.helper_functions import parse_expr, checkLazyNew
 
 
 def run(script, variables=None, optimize=True, stack=None):
-    tree = parse_expr(script)
+    tree = parse_expr.__wrapped__(script)
     if optimize:
         remove_dead_code(tree)
     out = io.StringIO()
@@ -109,5 +109,5 @@ def test_last_statement_of_a_block_stays():
 
 
 def test_unused_inputs_are_no_longer_requested():
-    tree = remove_dead_code(parse_expr("u = V1 * 2; V0"))
+    tree = remove_dead_code(parse_expr.__wrapped__("u = V1 * 2; V0"))
     assert checkLazyNew(tree, {"V0": None, "V1": None}, {}) == {"V0"}