nodes/rgthree-rgthree-comfy
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Power Puter: Add for loop, and augassign capabilities.

Browse Source
This commit is contained in:
rgthree
2025-06-23 00:15:18 -04:00
parent 51bc089334
commit f754c47658
3 changed files with 107 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/power_puter.py
+45 -11
View File
@@ -15,7 +15,7 @@ import re
import time
import operator as op
from typing import Any, Callable, Optional, Union
from typing import Any, Callable, Iterable, Optional, Union
from .constants import get_category, get_name
from .utils import ByPassTypeTuple, FlexibleOptionalInputType, any_type, get_dict_value
@@ -66,6 +66,7 @@ _FUNCTIONS = {
Function(name="re", call=re.compile, args=(1, 1)),
Function(name="len", call=len, args=(1, 1)),
Function(name="enumerate", call=enumerate, args=(1, 1)),
Function(name="range", call=range, args=(1, 3)),
# Casts
Function(name="int", call=int, args=(1, 1)),
Function(name="float", call=float, args=(1, 1)),
@@ -314,7 +315,7 @@ class _Puter:
log_node_warn(_NODE_NAME, f'No input node found for "{input_name}". ')
return None
def _eval_statement(self, stmt: ast.stmt, ctx: dict, prev_stmt: Union[ast.stmt, None] = None):
def _eval_statement(self, stmt: ast.AST, ctx: dict, prev_stmt: Union[ast.AST, None] = None):
"""Evaluates an ast.stmt."""
if '__returned__' in ctx:
@@ -351,7 +352,8 @@ class _Puter:
# Like: node(14)['inputs']['sampler_name'] (Subscript)
item = self._eval_statement(stmt.value, ctx=ctx)
attr = None
if hasattr(stmt, 'attr'):
# if hasattr(stmt, 'attr'):
if isinstance(stmt, ast.Attribute):
attr = stmt.attr
else:
# Slice could be a name or a constant; evaluate it
@@ -400,6 +402,20 @@ class _Puter:
return val
raise NameError(f"Name not found: {stmt.id}")
if isinstance(stmt, ast.For):
for_iter = self._eval_statement(stmt.iter, ctx=ctx)
for item in for_iter:
# Set the for var(s)
if isinstance(stmt.target, ast.Name):
ctx[stmt.target.id] = item
elif isinstance(stmt.target, ast.Tuple): # dict, like `for k, v in d.entries()`
for i, elt in enumerate(stmt.target.elts):
ctx[elt.id] = item[i]
bodies = stmt.body if isinstance(stmt.body, list) else [stmt.body]
for body in bodies:
value = self._eval_statement(body, ctx=ctx)
return None
if isinstance(stmt, ast.ListComp):
# Like: [v.lora for name, v in node(19).inputs.items() if name.startswith('lora_')]
# Like: [v.lower() for v in lora_list]
@@ -423,12 +439,16 @@ class _Puter:
else:
raise ValueError('Na')
gen_iters = None
# A call, like my_dct.items(), or a named ctx list
if isinstance(gen.iter, ast.Call):
gen_iters = self._eval_statement(gen.iter, ctx=gen_ctx)
elif isinstance(gen.iter, (ast.Name, ast.Attribute, ast.List, ast.Tuple)):
gen_iters = self._eval_statement(gen.iter, ctx=gen_ctx)
if not isinstance(gen_iters, Iterable):
raise ValueError('No iteraors found for list comprehension')
for gen_iter in gen_iters:
if_ctx = {**gen_ctx}
if isinstance(gen.target, ast.Tuple): # dict, like `for k, v in d.entries()`
@@ -518,16 +538,27 @@ class _Puter:
return value
# Assign a variable and add it to our ctx.
if isinstance(stmt, ast.Assign):
value = self._eval_statement(stmt.value, ctx=ctx)
if len(stmt.targets) != 1:
raise ValueError('Expected length of assign targets to be 1')
target = stmt.targets[0]
if isinstance(target, ast.Tuple):
if isinstance(stmt, (ast.Assign, ast.AugAssign)):
if isinstance(stmt, ast.AugAssign):
left = self._eval_statement(stmt.target, ctx=ctx)
right = self._eval_statement(stmt.value, ctx=ctx)
value = _OPERATORS[type(stmt.op)](left, right)
target = stmt.target
else:
value = self._eval_statement(stmt.value, ctx=ctx)
if len(stmt.targets) != 1:
raise ValueError('Expected length of assign targets to be 1')
target = stmt.targets[0]
if isinstance(target, ast.Tuple): # like `a, z = (1,2)` (ast.Assign only)
for i, elt in enumerate(target.elts):
ctx[elt.id] = value[i]
else: # Should be ast.Name
elif isinstance(target, ast.Name): # like `a = 1``
ctx[target.id] = value
elif isinstance(target, ast.Subscript) and isinstance(target.value, ast.Name): # like `a[0] = 1``
ctx[target.value.id][self._eval_statement(target.slice, ctx=ctx)] = value
else:
raise ValueError('Unhandled target type for Assign.')
return value
# For assigning a var in a list comprehension.
@@ -538,7 +569,10 @@ class _Puter:
return value
if isinstance(stmt, ast.Return):
value = self._eval_statement(stmt.value, ctx=ctx)
if stmt.value is None:
value = None
else:
value = self._eval_statement(stmt.value, ctx=ctx)
# Mark that we have a return value, as we may be deeper in evaluation, like going through an
# if condition's body.
ctx['__returned__'] = value
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "rgthree-comfy"
description = "Making ComfyUI more comfortable."
version = "1.0.2506222249"
version = "1.0.2506230015"
license = { file = "LICENSE" }
dependencies = []
src_web/comfyui/tests/power_puter.ts
+61 -7
View File
@@ -6,7 +6,15 @@ import {describe, should, beforeEach, expect, describeRun} from "../testing/runn
const env = new ComfyUITestEnvironment();
function setPowerPuterValue(node: LGraphNode, value: string, outputType: string = "STRING") {
function setPowerPuterValue(node: LGraphNode, outputType: string, value: string) {
// Strip as much whitespace on first non-empty line from all lines.
if (value.includes('\n')) {
value = value.replace(/^\n/gm, '')
const strip = value.match(/^(.*?)\S/)?.[1]?.length;
if (strip) {
value = value.replace(new RegExp(`^.{${strip}}`, 'mg'), '')
}
}
node.widgets![1]!.value = value;
node.widgets![0]!.value = outputType;
}
@@ -24,11 +32,11 @@ describe("TestPowerPuter", async () => {
});
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"],
const checks: Array<[string, string, string]> = [
["1", "1", "STRING"],
['"abc"', "abc", "STRING"],
["1 + 2", "3", "STRING"],
['"abc" + "xyz"', "abcxyz", "STRING"],
// INT
["1", "1", "INT"],
["1 + 2", "3", "INT"],
@@ -40,7 +48,7 @@ describe("TestPowerPuter", async () => {
["1 - 1", "False", "BOOLEAN"],
];
for (const data of checks) {
setPowerPuterValue(powerPuter, data[0], data[2]);
setPowerPuterValue(powerPuter, data[2], data[0]);
await env.queuePrompt();
expect(displayAny.widgets![0]!.value).toBe(data[0], data[1]);
}
@@ -53,4 +61,50 @@ describe("TestPowerPuter", async () => {
await should("handle complex inputs", async () => {
// TODO
});
await should("handle a for loop", async () => {
setPowerPuterValue(
powerPuter,
"STRING",
`
a = 0
b = ''
for n in range(4):
a += n
for m in range(2):
b += f'{str(n)}-{str(m)}.'
f'a:{a} b:{b}'
`,
);
await env.queuePrompt();
expect(displayAny.widgets![0]!.value).toBe("a:6 b:0-0.0-1.1-0.1-1.2-0.2-1.3-0.3-1.");
});
await should("handle assigning with a subscript slice", async () => {
setPowerPuterValue(
powerPuter,
"STRING",
`
a = [1,2,0]
a[a[2]] = 3
tuple(a)
`,
);
await env.queuePrompt();
expect(displayAny.widgets![0]!.value).toBe("(3, 2, 0)");
});
await should("handle aug assigning with a subscript slice", async () => {
setPowerPuterValue(
powerPuter,
"STRING",
`
a = [1,2,0]
a[a[2]] += 3
tuple(a)
`,
);
await env.queuePrompt();
expect(displayAny.widgets![0]!.value).toBe("(4, 2, 0)");
});
});