Power Puter: Add re() method for regex compilation; Allow searching for nodes names by regex; better recursive handling of multi-looped list comprehsion; breaking change: nodes() now returns a lit.

.gitignore

@@ -8,4 +8,6 @@ rgthree_config.json
 web/rgthree_config.js
 web/comfyui/rgthree_config.js
 userdata/
-userdata/**
+userdata/**
+web/comfyui/testing/
+web/comfyui/tests/

__build__.py

@@ -46,7 +46,7 @@ def log_step_info(msg:str, status='info'):
   step_infos.append({"msg": f'  - {msg}', "type": status})
 
 
-def build(with_tests = False, fix = False):
+def build(without_tests = True, fix = False):
 
   THIS_DIR = os.path.dirname(os.path.abspath(__file__))
   DIR_SRC_WEB = os.path.abspath(f'{THIS_DIR}/src_web/')
@@ -84,10 +84,10 @@ def build(with_tests = False, fix = False):
   checked = subprocess.run(["node", "./node_modules/typescript/bin/tsc"], check=True)
   log_step(status="Done")
 
-  if with_tests:
+  if not without_tests:
     log_step(msg='Removing directories (KEEPING TESTING)', status="Notice")
   else:
-    log_step(msg='Removing uneeded directories')
+    log_step(msg='Removing unneeded directories')
     test_path = os.path.join(DIR_WEB, 'comfyui', 'tests')
     if os.path.exists(test_path):
       rmtree(test_path)
@@ -137,10 +137,10 @@ def build(with_tests = False, fix = False):
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("-t", "--with-tests", default=False, action="store_true")
+  parser.add_argument("-t", "--no-tests", default=False, action="store_true")
   parser.add_argument("-f", "--fix", default=False, action="store_true")
   args = parser.parse_args()
 
   start = time.time()
-  build(with_tests=args.with_tests, fix=args.fix)
+  build(without_tests=args.no_tests, fix=args.fix)
   print(f'Finished all in {round(time.time() - start, 3)}s')

py/power_puter.py

@@ -17,6 +17,7 @@ from typing import Any, Callable
 import operator as op
 from .constants import get_category, get_name
 from .utils import FlexibleOptionalInputType, any_type, get_dict_value
+from .log import log_node_warn
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
@@ -44,6 +45,7 @@ _FUNCTIONS = {
     Function(name="max", call=max, args=(2, None)),
     Function(name="random_int", call=random.randint, args=(2, 2)),
     Function(name="random_choice", call=random.choice, args=(2, None)),
+    Function(name="re", call=re.compile, args=(1, 1)),
     # Casts
     Function(name="int", call=int, args=(1, 1)),
     Function(name="float", call=float, args=(1, 1)),
@@ -62,7 +64,6 @@ _FUNCTIONS = {
 # tell.
 _GLOBAL_FUNCTION_NAMES = ['node', 'nodes']
 
-
 # Special functions by class type (called from the Attrs.)
 _SPECIAL_FUNCTIONS = {
   "Power Lora Loader (rgthree)": {
@@ -79,7 +80,6 @@ _SPECIAL_FUNCTIONS = {
   }
 }
 
-
 _OPERATORS = {
   ast.Add: op.add,
   ast.Sub: op.sub,
@@ -102,11 +102,13 @@ _OPERATORS = {
 
 _IS_CHANGED_GLOBAL = 0
 
+_NODE_NAME = get_name("Power Puter")
+
 
 class RgthreePowerPuter:
   """A powerful node that can compute and evaluate expressions and output as various types."""
 
-  NAME = get_name("Power Puter")
+  NAME = _NODE_NAME
   CATEGORY = get_category()
 
   @classmethod
@@ -178,6 +180,9 @@ class _Puter:
     self._code = code
     self._workflow = workflow
     self._prompt = prompt
+    self._prompt_nodes = []
+    if self._prompt:
+      self._prompt_nodes = [{'id': id} | {**node} for id, node in self._prompt.items()]
 
   def execute(self, code=str | None) -> Any:
     """Evaluates a the code block."""
@@ -191,26 +196,33 @@ class _Puter:
         break
     return last_value
 
-  def _get_nodes(self, node_id: int | str | None = None) -> dict[str, Any]:
-    """Get a dict of the nodes that match the node_id, or all the nodes in the prompt."""
+  def _get_nodes(self, node_id: int | str | re.Pattern | None = None) -> list[Any]:
+    """Get a list of the nodes that match the node_id, or all the nodes in the prompt."""
+    log_node_warn(
+      _NODE_NAME,
+      "There was a breaking change to `nodes()` in Power Puter. The return is now a list of nodes,"
+      " rather than a dict. Update any Power Puter code."
+    )
+    nodes = self._prompt_nodes.copy()
     if not node_id:
-      return {**self._prompt}
-    node_id = str(node_id)
-    nodes = []
-    if re.match(r'\d+$', node_id):
-      nodes = {id: n for id, n in self._prompt.items() if node_id == id}
-    if not nodes:
-      nodes = {
-        id: n for id, n in self._prompt.items() if node_id == get_dict_value(n, '_meta.title', '')
-      }
-    return nodes
+      return nodes
 
-  def _get_node(self, node_id: int | str):
+    if isinstance(node_id, re.Pattern):
+      found = [n for n in nodes if re.search(node_id, get_dict_value(n, '_meta.title', ''))]
+    else:
+      node_id = str(node_id)
+      found = None
+      if re.match(r'\d+$', node_id):
+        found = [n for n in nodes if node_id == n['id']]
+      if not found:
+        found = [n for n in nodes if node_id == get_dict_value(n, '_meta.title', '')]
+    return found
+
+  def _get_node(self, node_id: int | str | re.Pattern) -> Any | None:
     """Returns a prompt-node from the hidden prompt."""
-    node_id = str(node_id)
-    nodes = [n for n in self._get_nodes(node_id).values()]
+    nodes = self._get_nodes(node_id)
     if nodes and len(nodes) > 1:
-      print('ERROR more than one node, returning first.')
+      log_node_warn(_NODE_NAME, f"More than one node found for '{node_id}'. Returning first.")
     return nodes[0] if nodes else None
 
   def _eval_statement(self, stmt: ast.stmt, ctx: dict | None = None):
@@ -285,42 +297,56 @@ class _Puter:
       # Like: [v.lower() for v in lora_list]
       # Like: [v for v in l if v.startswith('B')]
       # Like: [v.lower() for v in l if v.startswith('B') or v.startswith('F')]
+      # ---
+      # Like: [l for n in nodes(re('Loras')).values() if (l := n.loras)]
       final_list = []
 
-      for gen in stmt.generators:
-        gen_ctx = {**ctx}
-        if isinstance(gen.target, ast.Tuple):
-          gen_ctx[gen.target.elts[0].id] = None
-          gen_ctx[gen.target.elts[1].id] = None
-        elif isinstance(gen.target, ast.Name):
+      gen_ctx = {**ctx}
+
+      generators = [*stmt.generators]
+
+      def handle_gen(generators: list[ast.comprehension]):
+        gen = generators.pop(0)
+        if isinstance(gen.target, ast.Name):
           gen_ctx[gen.target.id] = None
+        elif isinstance(gen.target, ast.Tuple):  # dict, like `for k, v in d.entries()`
+          for elt in gen.target.elts:
+            gen_ctx[elt.id] = None
+        else:
+          raise ValueError('Na')
 
         # A call, like my_dct.items(), or a named ctx list
         if isinstance(gen.iter, ast.Call):
-          iter = self._eval_statement(gen.iter.func, ctx=gen_ctx)()
+          gen_iters = self._eval_statement(gen.iter.func, ctx=gen_ctx)()
         elif isinstance(gen.iter, (ast.Name, ast.Attribute)):
-          iter = self._eval_statement(gen.iter, ctx=gen_ctx)
+          gen_iters = self._eval_statement(gen.iter, ctx=gen_ctx)
 
-        for v in iter:
-          # Unpack if we were a dict, otherwise it's just the item
-          if isinstance(v, (tuple, list)):
-            gen_ctx[gen.target.elts[0].id] = v[0]
-            gen_ctx[gen.target.elts[1].id] = v[1]
+        for gen_iter in gen_iters:
+          if_ctx = {**gen_ctx}
+          if isinstance(gen.target, ast.Tuple):  # dict, like `for k, v in d.entries()`
+            for i, elt in enumerate(gen.target.elts):
+              if_ctx[elt.id] = gen_iter[i]
           else:
-            gen_ctx[gen.target.id] = v
-
+            if_ctx[gen.target.id] = gen_iter
           good = True
           for ifcall in gen.ifs:
-            if not self._eval_statement(ifcall, ctx=gen_ctx):
+            if not self._eval_statement(ifcall, ctx=if_ctx):
               good = False
               break
           if not good:
             continue
+          gen_ctx.update(if_ctx)
+          if len(generators):
+            handle_gen(generators)
+          else:
+            final_list.append(self._eval_statement(stmt.elt, gen_ctx))
+        generators.insert(0, gen)
 
-          final_list.append(self._eval_statement(stmt.elt, gen_ctx))
-        return final_list
+      handle_gen(generators)
+      return final_list
 
     if isinstance(stmt, ast.Call):
+      call = None
       if isinstance(stmt.func, ast.Attribute):
         call = self._eval_statement(stmt.func, ctx=ctx)
       if isinstance(stmt.func, ast.Name):
@@ -335,6 +361,8 @@ class _Puter:
             toErr = " or more" if fn.args[1] is None else f" to {fn.args[1]}"
             raise SyntaxError(f"Invalid function call: {fn.name} requires {fn.args[0]}{toErr} args")
       if not call:
+        # print(type(stmt))
+        # print(type(stmt.func))
         raise ValueError(f'No call for ast.Call {stmt}')
       args = []
       for arg in stmt.args:
@@ -363,7 +391,14 @@ class _Puter:
     # Assign a variable and add it to our ctx.
     if isinstance(stmt, ast.Assign):
       value = self._eval_statement(stmt.value, ctx=ctx)
-      self._ctx[stmt.targets[0].id] = value
+      ctx[stmt.targets[0].id] = value
+      return value
+
+    # For assigning a var in a list comprehension.
+    # Like [name for node in node_list if (name := node.name)]
+    if isinstance(stmt, ast.NamedExpr):
+      value = self._eval_statement(stmt.value, ctx=ctx)
+      ctx[stmt.target.id] = value
       return value
 
     raise TypeError(stmt)

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "rgthree-comfy"
 description = "Making ComfyUI more comfortable."
-version = "1.0.2505182359"
+version = "1.0.2505212110"
 license = { file = "LICENSE" }
 dependencies = []
 

src_web/comfyui/tests/power_puter.ts

@@ -0,0 +1,56 @@
+import type {LGraphNode} from "@comfyorg/litegraph";
+
+import {NodeTypesString} from "../constants";
+import {ComfyUITestEnvironment} from "../testing/comfyui_env";
+import {describe, should, beforeEach, expect, describeRun} from "../testing/runner.js";
+
+const env = new ComfyUITestEnvironment();
+
+function setPowerPuterValue(node: LGraphNode, value: string, outputType: string = "STRING") {
+  node.widgets![1]!.value = value;
+  node.widgets![0]!.value = outputType;
+}
+
+describe("TestPowerPuter", async () => {
+  let powerPuter!: LGraphNode;
+  let displayAny!: LGraphNode;
+
+  await beforeEach(async () => {
+    await env.clear();
+    powerPuter = await env.addNode(NodeTypesString.POWER_PUTER);
+    displayAny = await env.addNode(NodeTypesString.DISPLAY_ANY);
+    powerPuter.connect(0, displayAny, 0);
+    await env.wait();
+  });
+
+  await should("output constants and concatenation", async () => {
+    const checks: Array<[string, string] | [string, string, string]> = [
+      ["1", "1"],
+      ['"abc"', "abc"],
+      ["1 + 2", "3"],
+      ['"abc" + "xyz"', "abcxyz"],
+      // INT
+      ["1", "1", "INT"],
+      ["1 + 2", "3", "INT"],
+      // FLOAT
+      ["1", "1.0", "FLOAT"],
+      ["1.3 + 2.8", "4.1", "FLOAT"],
+      // BOOL
+      ["1", "True", "BOOL"],
+      ["1 - 1", "False", "BOOL"],
+    ];
+    for (const data of checks) {
+      setPowerPuterValue(powerPuter, data[0], data[2]);
+      await env.queuePrompt();
+      expect(displayAny.widgets![0]!.value).toBe(data[0], data[1]);
+    }
+  });
+
+  await should("handle inputs", async () => {
+    // TODO
+  });
+
+  await should("handle complex inputs", async () => {
+    // TODO
+  });
+});