py2dag 0.1.14__py3-none-any.whl → 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- py2dag/export_dagre.py +15 -2
- py2dag/parser.py +346 -176
- {py2dag-0.1.14.dist-info → py2dag-0.2.0.dist-info}/METADATA +1 -1
- py2dag-0.2.0.dist-info/RECORD +11 -0
- py2dag-0.1.14.dist-info/RECORD +0 -11
- {py2dag-0.1.14.dist-info → py2dag-0.2.0.dist-info}/LICENSE +0 -0
- {py2dag-0.1.14.dist-info → py2dag-0.2.0.dist-info}/WHEEL +0 -0
- {py2dag-0.1.14.dist-info → py2dag-0.2.0.dist-info}/entry_points.txt +0 -0
py2dag/export_dagre.py
CHANGED
@@ -48,9 +48,22 @@ HTML_TEMPLATE = """<!doctype html>
|
|
48
48
|
// Ensure edges have an object for labels/attrs to avoid TypeErrors
|
49
49
|
g.setDefaultEdgeLabel(() => ({}));
|
50
50
|
|
51
|
-
// Add op nodes
|
51
|
+
// Add op nodes with basic styling for control nodes
|
52
52
|
(plan.ops || []).forEach(op => {
|
53
|
-
|
53
|
+
let label = op.op;
|
54
|
+
let klass = 'op';
|
55
|
+
if (op.op === 'COND.eval') {
|
56
|
+
const kind = (op.args && op.args.kind) || 'if';
|
57
|
+
label = (kind.toUpperCase()) + ' ' + (op.args && op.args.expr ? op.args.expr : '');
|
58
|
+
klass = 'note';
|
59
|
+
} else if (op.op === 'ITER.eval') {
|
60
|
+
label = 'FOR ' + (op.args && op.args.expr ? op.args.expr : '');
|
61
|
+
klass = 'note';
|
62
|
+
} else if (op.op === 'PHI') {
|
63
|
+
label = 'PHI' + (op.args && op.args.var ? ` (${op.args.var})` : '');
|
64
|
+
klass = 'note';
|
65
|
+
}
|
66
|
+
g.setNode(op.id, { label, class: klass, padding: 8 });
|
54
67
|
});
|
55
68
|
|
56
69
|
// Add output nodes and edges from source to output
|
py2dag/parser.py
CHANGED
@@ -1,7 +1,7 @@
|
|
1
1
|
import ast
|
2
2
|
import json
|
3
3
|
import re
|
4
|
-
from typing import Any, Dict, List, Optional
|
4
|
+
from typing import Any, Dict, List, Optional, Tuple, Set
|
5
5
|
|
6
6
|
VALID_NAME_RE = re.compile(r'^[a-z_][a-z0-9_]{0,63}$')
|
7
7
|
|
@@ -41,184 +41,243 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
|
|
41
41
|
module = ast.parse(source)
|
42
42
|
|
43
43
|
def _parse_fn(fn: ast.AST) -> Dict[str, Any]:
|
44
|
-
defined: set[str] = set()
|
45
44
|
ops: List[Dict[str, Any]] = []
|
46
45
|
outputs: List[Dict[str, str]] = []
|
47
46
|
settings: Dict[str, Any] = {}
|
48
47
|
|
49
48
|
returned_var: Optional[str] = None
|
50
|
-
|
49
|
+
|
50
|
+
# Enforce no-args top-level function signature
|
51
51
|
try:
|
52
|
-
|
53
|
-
|
54
|
-
|
52
|
+
fargs = getattr(fn, "args") # type: ignore[attr-defined]
|
53
|
+
has_params = bool(
|
54
|
+
getattr(fargs, "posonlyargs", []) or fargs.args or fargs.vararg or fargs.kwonlyargs or fargs.kwarg
|
55
|
+
)
|
56
|
+
if has_params:
|
57
|
+
raise DSLParseError("Top-level function must not accept parameters")
|
58
|
+
except AttributeError:
|
55
59
|
pass
|
56
60
|
|
57
|
-
|
61
|
+
# SSA state
|
62
|
+
versions: Dict[str, int] = {}
|
63
|
+
latest: Dict[str, str] = {}
|
64
|
+
context_suffix: str = ""
|
65
|
+
ctx_counts: Dict[str, int] = {"if": 0, "loop": 0, "while": 0}
|
66
|
+
|
67
|
+
def _ssa_new(name: str) -> str:
|
68
|
+
if not VALID_NAME_RE.match(name):
|
69
|
+
raise DSLParseError(f"Invalid variable name: {name}")
|
70
|
+
versions[name] = versions.get(name, 0) + 1
|
71
|
+
base = f"{name}_{versions[name]}"
|
72
|
+
ssa = f"{base}@{context_suffix}" if context_suffix else base
|
73
|
+
latest[name] = ssa
|
74
|
+
return ssa
|
75
|
+
|
76
|
+
def _ssa_get(name: str) -> str:
|
77
|
+
if name not in latest:
|
78
|
+
raise DSLParseError(f"Undefined dependency: {name}")
|
79
|
+
return latest[name]
|
80
|
+
|
81
|
+
def _collect_name_loads(node: ast.AST) -> List[str]:
|
58
82
|
names: List[str] = []
|
59
83
|
for n in ast.walk(node):
|
60
84
|
if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load):
|
61
85
|
if n.id not in names:
|
62
86
|
names.append(n.id)
|
63
87
|
return names
|
64
|
-
|
65
|
-
|
88
|
+
|
89
|
+
def _collect_value_deps(node: ast.AST) -> List[str]:
|
90
|
+
"""Collect variable name dependencies from an expression, excluding callee names in Call.func.
|
91
|
+
|
92
|
+
For example, for range(n) -> ['n'] (not 'range'). For cond(a) -> ['a'] (not 'cond').
|
93
|
+
For obj.attr -> ['obj'].
|
94
|
+
"""
|
95
|
+
callees: set[str] = set()
|
96
|
+
|
97
|
+
def mark_callee(func: ast.AST):
|
98
|
+
for n in ast.walk(func):
|
99
|
+
if isinstance(n, ast.Name):
|
100
|
+
callees.add(n.id)
|
101
|
+
|
102
|
+
# First collect callee name ids appearing under Call.func
|
103
|
+
for n in ast.walk(node):
|
104
|
+
if isinstance(n, ast.Call):
|
105
|
+
mark_callee(n.func)
|
106
|
+
|
107
|
+
# Then collect normal loads and drop any that are marked as callees
|
108
|
+
deps: List[str] = []
|
109
|
+
for n in ast.walk(node):
|
110
|
+
if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load):
|
111
|
+
if n.id not in callees and n.id not in deps:
|
112
|
+
deps.append(n.id)
|
113
|
+
return deps
|
114
|
+
|
115
|
+
def _stringify(node: ast.AST) -> str:
|
116
|
+
try:
|
117
|
+
return ast.unparse(node) # type: ignore[attr-defined]
|
118
|
+
except Exception:
|
119
|
+
return node.__class__.__name__
|
120
|
+
|
121
|
+
def _emit_assign_from_call(var_name: str, call: ast.Call) -> str:
|
122
|
+
op_name = _get_call_name(call.func)
|
123
|
+
deps: List[str] = []
|
124
|
+
|
125
|
+
def _expand_star_name(ssa_var: str) -> List[str]:
|
126
|
+
# Expand if previous op was a PACK.*
|
127
|
+
for prev in reversed(ops):
|
128
|
+
if prev.get("id") == ssa_var and prev.get("op") in {"PACK.list", "PACK.tuple"}:
|
129
|
+
return list(prev.get("deps", []))
|
130
|
+
return [ssa_var]
|
131
|
+
|
132
|
+
for arg in call.args:
|
133
|
+
if isinstance(arg, ast.Starred):
|
134
|
+
star_val = arg.value
|
135
|
+
if isinstance(star_val, ast.Name):
|
136
|
+
deps.extend(_expand_star_name(_ssa_get(star_val.id)))
|
137
|
+
elif isinstance(star_val, (ast.List, ast.Tuple)):
|
138
|
+
for elt in star_val.elts:
|
139
|
+
if not isinstance(elt, ast.Name):
|
140
|
+
raise DSLParseError("Starred list/tuple elements must be names")
|
141
|
+
deps.append(_ssa_get(elt.id))
|
142
|
+
else:
|
143
|
+
raise DSLParseError("*args must be a name or list/tuple of names")
|
144
|
+
elif isinstance(arg, ast.Name):
|
145
|
+
deps.append(_ssa_get(arg.id))
|
146
|
+
elif isinstance(arg, (ast.List, ast.Tuple)):
|
147
|
+
for elt in arg.elts:
|
148
|
+
if not isinstance(elt, ast.Name):
|
149
|
+
raise DSLParseError("List/Tuple positional args must be variable names")
|
150
|
+
deps.append(_ssa_get(elt.id))
|
151
|
+
else:
|
152
|
+
raise DSLParseError("Positional args must be variable names or lists/tuples of names")
|
153
|
+
|
154
|
+
kwargs: Dict[str, Any] = {}
|
155
|
+
for kw in call.keywords:
|
156
|
+
if kw.arg is None:
|
157
|
+
v = kw.value
|
158
|
+
if isinstance(v, ast.Dict):
|
159
|
+
lit = _literal(v)
|
160
|
+
for k, val in lit.items():
|
161
|
+
kwargs[str(k)] = val
|
162
|
+
elif isinstance(v, ast.Name):
|
163
|
+
deps.append(_ssa_get(v.id))
|
164
|
+
else:
|
165
|
+
raise DSLParseError("**kwargs must be a dict literal or a variable name")
|
166
|
+
else:
|
167
|
+
if isinstance(kw.value, ast.Name):
|
168
|
+
deps.append(_ssa_get(kw.value.id))
|
169
|
+
else:
|
170
|
+
kwargs[kw.arg] = _literal(kw.value)
|
171
|
+
|
172
|
+
ssa = _ssa_new(var_name)
|
173
|
+
ops.append({"id": ssa, "op": op_name, "deps": deps, "args": kwargs})
|
174
|
+
return ssa
|
175
|
+
|
176
|
+
def _emit_assign_from_fstring(var_name: str, fstr: ast.JoinedStr) -> str:
|
177
|
+
deps: List[str] = []
|
178
|
+
parts: List[str] = []
|
179
|
+
for item in fstr.values:
|
180
|
+
if isinstance(item, ast.Constant) and isinstance(item.value, str):
|
181
|
+
parts.append(item.value)
|
182
|
+
elif isinstance(item, ast.FormattedValue) and isinstance(item.value, ast.Name):
|
183
|
+
deps.append(_ssa_get(item.value.id))
|
184
|
+
parts.append("{" + str(len(deps) - 1) + "}")
|
185
|
+
else:
|
186
|
+
raise DSLParseError("f-strings may only contain variable names")
|
187
|
+
template = "".join(parts)
|
188
|
+
ssa = _ssa_new(var_name)
|
189
|
+
ops.append({
|
190
|
+
"id": ssa,
|
191
|
+
"op": "TEXT.format",
|
192
|
+
"deps": deps,
|
193
|
+
"args": {"template": template},
|
194
|
+
})
|
195
|
+
return ssa
|
196
|
+
|
197
|
+
def _emit_assign_from_literal_or_pack(var_name: str, value: ast.AST) -> str:
|
198
|
+
try:
|
199
|
+
lit = _literal(value)
|
200
|
+
ssa = _ssa_new(var_name)
|
201
|
+
ops.append({
|
202
|
+
"id": ssa,
|
203
|
+
"op": "CONST.value",
|
204
|
+
"deps": [],
|
205
|
+
"args": {"value": lit},
|
206
|
+
})
|
207
|
+
return ssa
|
208
|
+
except DSLParseError:
|
209
|
+
if isinstance(value, (ast.List, ast.Tuple)):
|
210
|
+
elts = value.elts
|
211
|
+
deps: List[str] = []
|
212
|
+
for elt in elts:
|
213
|
+
if not isinstance(elt, ast.Name):
|
214
|
+
raise DSLParseError("Only names allowed in non-literal list/tuple assignment")
|
215
|
+
deps.append(_ssa_get(elt.id))
|
216
|
+
kind = "list" if isinstance(value, ast.List) else "tuple"
|
217
|
+
ssa = _ssa_new(var_name)
|
218
|
+
ops.append({
|
219
|
+
"id": ssa,
|
220
|
+
"op": f"PACK.{kind}",
|
221
|
+
"deps": deps,
|
222
|
+
"args": {},
|
223
|
+
})
|
224
|
+
return ssa
|
225
|
+
raise
|
226
|
+
|
227
|
+
def _emit_assign_from_comp(var_name: str, node: ast.AST) -> str:
|
228
|
+
name_deps = [n for n in _collect_name_loads(node) if n in latest]
|
229
|
+
for n in name_deps:
|
230
|
+
if n not in latest:
|
231
|
+
raise DSLParseError(f"Undefined dependency: {n}")
|
232
|
+
kind = (
|
233
|
+
"listcomp" if isinstance(node, ast.ListComp) else
|
234
|
+
"setcomp" if isinstance(node, ast.SetComp) else
|
235
|
+
"dictcomp" if isinstance(node, ast.DictComp) else
|
236
|
+
"genexpr"
|
237
|
+
)
|
238
|
+
deps = [_ssa_get(n) for n in name_deps]
|
239
|
+
ssa = _ssa_new(var_name)
|
240
|
+
ops.append({
|
241
|
+
"id": ssa,
|
242
|
+
"op": f"COMP.{kind}",
|
243
|
+
"deps": deps,
|
244
|
+
"args": {},
|
245
|
+
})
|
246
|
+
return ssa
|
247
|
+
|
248
|
+
def _emit_cond(node: ast.AST, kind: str = "if") -> str:
|
249
|
+
expr = _stringify(node)
|
250
|
+
deps = [_ssa_get(n) for n in _collect_value_deps(node)]
|
251
|
+
ssa = _ssa_new("cond")
|
252
|
+
ops.append({"id": ssa, "op": "COND.eval", "deps": deps, "args": {"expr": expr, "kind": kind}})
|
253
|
+
return ssa
|
254
|
+
|
255
|
+
def _emit_iter(node: ast.AST) -> str:
|
256
|
+
expr = _stringify(node)
|
257
|
+
deps = [_ssa_get(n) for n in _collect_value_deps(node)]
|
258
|
+
ssa = _ssa_new("iter")
|
259
|
+
ops.append({"id": ssa, "op": "ITER.eval", "deps": deps, "args": {"expr": expr, "kind": "for"}})
|
260
|
+
return ssa
|
261
|
+
|
262
|
+
def _parse_stmt(stmt: ast.stmt) -> Optional[str]:
|
263
|
+
nonlocal returned_var, versions, latest, context_suffix
|
66
264
|
if isinstance(stmt, ast.Assign):
|
67
265
|
if len(stmt.targets) != 1 or not isinstance(stmt.targets[0], ast.Name):
|
68
266
|
raise DSLParseError("Assignment targets must be simple names")
|
69
267
|
var_name = stmt.targets[0].id
|
70
|
-
if not VALID_NAME_RE.match(var_name):
|
71
|
-
raise DSLParseError(f"Invalid variable name: {var_name}")
|
72
|
-
if var_name in defined:
|
73
|
-
raise DSLParseError(f"Duplicate variable name: {var_name}")
|
74
|
-
|
75
268
|
value = stmt.value
|
76
269
|
if isinstance(value, ast.Await):
|
77
270
|
value = value.value
|
78
271
|
if isinstance(value, ast.Call):
|
79
|
-
|
80
|
-
|
81
|
-
deps: List[str] = []
|
82
|
-
|
83
|
-
def _expand_star_name(varname: str) -> List[str]:
|
84
|
-
# Try to expand a previously packed list/tuple variable into its element deps
|
85
|
-
for prev in reversed(ops):
|
86
|
-
if prev.get("id") == varname:
|
87
|
-
if prev.get("op") in {"PACK.list", "PACK.tuple"}:
|
88
|
-
return list(prev.get("deps", []))
|
89
|
-
break
|
90
|
-
return [varname]
|
91
|
-
for arg in value.args:
|
92
|
-
if isinstance(arg, ast.Starred):
|
93
|
-
star_val = arg.value
|
94
|
-
if isinstance(star_val, ast.Name):
|
95
|
-
if star_val.id not in defined:
|
96
|
-
raise DSLParseError(f"Undefined dependency: {star_val.id}")
|
97
|
-
deps.extend(_expand_star_name(star_val.id))
|
98
|
-
elif isinstance(star_val, (ast.List, ast.Tuple)):
|
99
|
-
for elt in star_val.elts:
|
100
|
-
if not isinstance(elt, ast.Name):
|
101
|
-
raise DSLParseError("Starred list/tuple elements must be names")
|
102
|
-
if elt.id not in defined:
|
103
|
-
raise DSLParseError(f"Undefined dependency: {elt.id}")
|
104
|
-
deps.append(elt.id)
|
105
|
-
else:
|
106
|
-
raise DSLParseError("*args must be a name or list/tuple of names")
|
107
|
-
elif isinstance(arg, ast.Name):
|
108
|
-
if arg.id not in defined:
|
109
|
-
raise DSLParseError(f"Undefined dependency: {arg.id}")
|
110
|
-
deps.append(arg.id)
|
111
|
-
elif isinstance(arg, (ast.List, ast.Tuple)):
|
112
|
-
for elt in arg.elts:
|
113
|
-
if not isinstance(elt, ast.Name):
|
114
|
-
raise DSLParseError("List/Tuple positional args must be variable names")
|
115
|
-
if elt.id not in defined:
|
116
|
-
raise DSLParseError(f"Undefined dependency: {elt.id}")
|
117
|
-
deps.append(elt.id)
|
118
|
-
else:
|
119
|
-
raise DSLParseError("Positional args must be variable names or lists/tuples of names")
|
120
|
-
|
121
|
-
kwargs: Dict[str, Any] = {}
|
122
|
-
for kw in value.keywords:
|
123
|
-
if kw.arg is None:
|
124
|
-
# **kwargs support: allow dict literal merge, or variable name as dep
|
125
|
-
v = kw.value
|
126
|
-
if isinstance(v, ast.Dict):
|
127
|
-
# Merge literal kwargs
|
128
|
-
lit = _literal(v)
|
129
|
-
for k, val in lit.items():
|
130
|
-
kwargs[str(k)] = val
|
131
|
-
elif isinstance(v, ast.Name):
|
132
|
-
if v.id not in defined:
|
133
|
-
raise DSLParseError(f"Undefined dependency: {v.id}")
|
134
|
-
deps.append(v.id)
|
135
|
-
else:
|
136
|
-
raise DSLParseError("**kwargs must be a dict literal or a variable name")
|
137
|
-
else:
|
138
|
-
# Support variable-name keyword args as dependencies; literals remain in args
|
139
|
-
if isinstance(kw.value, ast.Name):
|
140
|
-
name = kw.value.id
|
141
|
-
if name not in defined:
|
142
|
-
raise DSLParseError(f"Undefined dependency: {name}")
|
143
|
-
deps.append(name)
|
144
|
-
else:
|
145
|
-
kwargs[kw.arg] = _literal(kw.value)
|
146
|
-
|
147
|
-
ops.append({"id": var_name, "op": op_name, "deps": deps, "args": kwargs})
|
272
|
+
return _emit_assign_from_call(var_name, value)
|
148
273
|
elif isinstance(value, ast.JoinedStr):
|
149
|
-
|
150
|
-
deps: List[str] = []
|
151
|
-
parts: List[str] = []
|
152
|
-
for item in value.values:
|
153
|
-
if isinstance(item, ast.Constant) and isinstance(item.value, str):
|
154
|
-
parts.append(item.value)
|
155
|
-
elif isinstance(item, ast.FormattedValue) and isinstance(item.value, ast.Name):
|
156
|
-
name = item.value.id
|
157
|
-
if name not in defined:
|
158
|
-
raise DSLParseError(f"Undefined dependency: {name}")
|
159
|
-
deps.append(name)
|
160
|
-
parts.append("{" + str(len(deps) - 1) + "}")
|
161
|
-
else:
|
162
|
-
raise DSLParseError("f-strings may only contain variable names")
|
163
|
-
template = "".join(parts)
|
164
|
-
ops.append({
|
165
|
-
"id": var_name,
|
166
|
-
"op": "TEXT.format",
|
167
|
-
"deps": deps,
|
168
|
-
"args": {"template": template},
|
169
|
-
})
|
274
|
+
return _emit_assign_from_fstring(var_name, value)
|
170
275
|
elif isinstance(value, (ast.Constant, ast.List, ast.Tuple, ast.Dict)):
|
171
|
-
|
172
|
-
try:
|
173
|
-
lit = _literal(value)
|
174
|
-
ops.append({
|
175
|
-
"id": var_name,
|
176
|
-
"op": "CONST.value",
|
177
|
-
"deps": [],
|
178
|
-
"args": {"value": lit},
|
179
|
-
})
|
180
|
-
except DSLParseError:
|
181
|
-
if isinstance(value, (ast.List, ast.Tuple)):
|
182
|
-
elts = value.elts
|
183
|
-
names: List[str] = []
|
184
|
-
for elt in elts:
|
185
|
-
if not isinstance(elt, ast.Name):
|
186
|
-
raise DSLParseError("Only names allowed in non-literal list/tuple assignment")
|
187
|
-
if elt.id not in defined:
|
188
|
-
raise DSLParseError(f"Undefined dependency: {elt.id}")
|
189
|
-
names.append(elt.id)
|
190
|
-
kind = "list" if isinstance(value, ast.List) else "tuple"
|
191
|
-
ops.append({
|
192
|
-
"id": var_name,
|
193
|
-
"op": f"PACK.{kind}",
|
194
|
-
"deps": names,
|
195
|
-
"args": {},
|
196
|
-
})
|
197
|
-
else:
|
198
|
-
raise
|
276
|
+
return _emit_assign_from_literal_or_pack(var_name, value)
|
199
277
|
elif isinstance(value, (ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp)):
|
200
|
-
|
201
|
-
name_deps = [n for n in _collect_name_deps(value) if n in defined]
|
202
|
-
# Ensure no undefined names used
|
203
|
-
for n in name_deps:
|
204
|
-
if n not in defined:
|
205
|
-
raise DSLParseError(f"Undefined dependency: {n}")
|
206
|
-
kind = (
|
207
|
-
"listcomp" if isinstance(value, ast.ListComp) else
|
208
|
-
"setcomp" if isinstance(value, ast.SetComp) else
|
209
|
-
"dictcomp" if isinstance(value, ast.DictComp) else
|
210
|
-
"genexpr"
|
211
|
-
)
|
212
|
-
ops.append({
|
213
|
-
"id": var_name,
|
214
|
-
"op": f"COMP.{kind}",
|
215
|
-
"deps": name_deps,
|
216
|
-
"args": {},
|
217
|
-
})
|
278
|
+
return _emit_assign_from_comp(var_name, value)
|
218
279
|
else:
|
219
280
|
raise DSLParseError("Right hand side must be a call or f-string")
|
220
|
-
defined.add(var_name)
|
221
|
-
|
222
281
|
elif isinstance(stmt, ast.Expr):
|
223
282
|
call = stmt.value
|
224
283
|
if isinstance(call, ast.Await):
|
@@ -237,8 +296,7 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
|
|
237
296
|
if len(call.args) != 1 or not isinstance(call.args[0], ast.Name):
|
238
297
|
raise DSLParseError("output requires a single variable name argument")
|
239
298
|
var = call.args[0].id
|
240
|
-
|
241
|
-
raise DSLParseError(f"Undefined output variable: {var}")
|
299
|
+
ssa_from = _ssa_get(var)
|
242
300
|
filename = None
|
243
301
|
for kw in call.keywords:
|
244
302
|
if kw.arg in {"as", "as_"}:
|
@@ -247,26 +305,16 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
|
|
247
305
|
raise DSLParseError("output only accepts 'as' keyword")
|
248
306
|
if filename is None or not isinstance(filename, str):
|
249
307
|
raise DSLParseError("output requires as=\"filename\"")
|
250
|
-
outputs.append({"from":
|
308
|
+
outputs.append({"from": ssa_from, "as": filename})
|
251
309
|
else:
|
252
310
|
raise DSLParseError("Only settings() and output() calls allowed as expressions")
|
311
|
+
return None
|
253
312
|
elif isinstance(stmt, ast.Return):
|
254
|
-
if i != len(fn.body) - 1: # type: ignore[index]
|
255
|
-
raise DSLParseError("return must be the last statement")
|
256
313
|
if isinstance(stmt.value, ast.Name):
|
257
|
-
|
258
|
-
if var not in defined:
|
259
|
-
raise DSLParseError(f"Undefined return variable: {var}")
|
260
|
-
returned_var = var
|
314
|
+
returned_var = _ssa_get(stmt.value.id)
|
261
315
|
elif isinstance(stmt.value, (ast.Constant, ast.List, ast.Tuple, ast.Dict)):
|
262
|
-
# Support returning a JSON-serialisable literal (str/num/bool/None, list/tuple, dict)
|
263
316
|
lit = _literal(stmt.value)
|
264
|
-
|
265
|
-
const_id = const_id_base
|
266
|
-
n = 1
|
267
|
-
while const_id in defined:
|
268
|
-
const_id = f"{const_id_base}_{n}"
|
269
|
-
n += 1
|
317
|
+
const_id = _ssa_new("return_value")
|
270
318
|
ops.append({
|
271
319
|
"id": const_id,
|
272
320
|
"op": "CONST.value",
|
@@ -276,25 +324,147 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
|
|
276
324
|
returned_var = const_id
|
277
325
|
else:
|
278
326
|
raise DSLParseError("return must return a variable name or literal")
|
279
|
-
|
280
|
-
|
281
|
-
|
327
|
+
return None
|
328
|
+
elif isinstance(stmt, ast.If):
|
329
|
+
# Evaluate condition
|
330
|
+
cond_id = _emit_cond(stmt.test, kind="if")
|
331
|
+
# Save pre-branch state
|
332
|
+
pre_versions = dict(versions)
|
333
|
+
pre_latest = dict(latest)
|
334
|
+
|
335
|
+
# THEN branch
|
336
|
+
then_ops_start = len(ops)
|
337
|
+
versions_then = dict(pre_versions)
|
338
|
+
latest_then = dict(pre_latest)
|
339
|
+
# Run then body with local state and context
|
340
|
+
saved_versions, saved_latest = versions, latest
|
341
|
+
saved_ctx = context_suffix
|
342
|
+
ctx_counts["if"] += 1
|
343
|
+
context_suffix = f"then{ctx_counts['if']}"
|
344
|
+
versions, latest = versions_then, latest_then
|
345
|
+
for inner in stmt.body:
|
346
|
+
_parse_stmt(inner)
|
347
|
+
versions_then, latest_then = versions, latest
|
348
|
+
versions, latest = saved_versions, saved_latest
|
349
|
+
context_suffix = saved_ctx
|
350
|
+
|
351
|
+
# ELSE branch
|
352
|
+
else_ops_start = len(ops)
|
353
|
+
versions_else = dict(pre_versions)
|
354
|
+
latest_else = dict(pre_latest)
|
355
|
+
saved_versions, saved_latest = versions, latest
|
356
|
+
saved_ctx = context_suffix
|
357
|
+
context_suffix = f"else{ctx_counts['if']}"
|
358
|
+
versions, latest = versions_else, latest_else
|
359
|
+
for inner in stmt.orelse or []:
|
360
|
+
_parse_stmt(inner)
|
361
|
+
versions_else, latest_else = versions, latest
|
362
|
+
versions, latest = saved_versions, saved_latest
|
363
|
+
context_suffix = saved_ctx
|
364
|
+
|
365
|
+
# Add cond dep to first op in each branch, if any
|
366
|
+
if len(ops) > then_ops_start:
|
367
|
+
ops[then_ops_start]["deps"] = [*ops[then_ops_start].get("deps", []), cond_id]
|
368
|
+
if len(ops) > else_ops_start:
|
369
|
+
ops[else_ops_start]["deps"] = [*ops[else_ops_start].get("deps", []), cond_id]
|
370
|
+
|
371
|
+
# Determine variables assigned in branches
|
372
|
+
then_assigned = {k for k in latest_then if pre_latest.get(k) != latest_then.get(k)}
|
373
|
+
else_assigned = {k for k in latest_else if pre_latest.get(k) != latest_else.get(k)}
|
374
|
+
all_assigned = then_assigned | else_assigned
|
375
|
+
for var in sorted(all_assigned):
|
376
|
+
left = latest_then.get(var, pre_latest.get(var))
|
377
|
+
right = latest_else.get(var, pre_latest.get(var))
|
378
|
+
if left is None or right is None:
|
379
|
+
# Variable does not exist pre-branch on one side; skip making it available post-merge
|
380
|
+
continue
|
381
|
+
phi_id = _ssa_new(var)
|
382
|
+
ops.append({"id": phi_id, "op": "PHI", "deps": [left, right], "args": {"var": var}})
|
383
|
+
return None
|
384
|
+
elif isinstance(stmt, (ast.For, ast.AsyncFor)):
|
385
|
+
# ITER over iterable
|
386
|
+
iter_id = _emit_iter(stmt.iter)
|
387
|
+
# Save pre-loop state
|
388
|
+
pre_versions = dict(versions)
|
389
|
+
pre_latest = dict(latest)
|
390
|
+
# Body state copy
|
391
|
+
body_ops_start = len(ops)
|
392
|
+
versions_body = dict(pre_versions)
|
393
|
+
latest_body = dict(pre_latest)
|
394
|
+
saved_versions, saved_latest = versions, latest
|
395
|
+
saved_ctx = context_suffix
|
396
|
+
ctx_counts["loop"] += 1
|
397
|
+
context_suffix = f"loop{ctx_counts['loop']}"
|
398
|
+
versions, latest = versions_body, latest_body
|
399
|
+
for inner in stmt.body:
|
400
|
+
_parse_stmt(inner)
|
401
|
+
versions_body, latest_body = versions, latest
|
402
|
+
versions, latest = saved_versions, saved_latest
|
403
|
+
context_suffix = saved_ctx
|
404
|
+
# Add iter dep to first op in body
|
405
|
+
if len(ops) > body_ops_start:
|
406
|
+
ops[body_ops_start]["deps"] = [*ops[body_ops_start].get("deps", []), iter_id]
|
407
|
+
# Loop-carried vars: only those existing pre-loop and reassigned in body
|
408
|
+
changed = {k for k in latest_body if pre_latest.get(k) != latest_body.get(k)}
|
409
|
+
carried = [k for k in changed if k in pre_latest]
|
410
|
+
for var in sorted(carried):
|
411
|
+
phi_id = _ssa_new(var)
|
412
|
+
ops.append({
|
413
|
+
"id": phi_id,
|
414
|
+
"op": "PHI",
|
415
|
+
"deps": [pre_latest[var], latest_body[var]],
|
416
|
+
"args": {"var": var},
|
417
|
+
})
|
418
|
+
return None
|
419
|
+
elif isinstance(stmt, ast.While):
|
420
|
+
cond_id = _emit_cond(stmt.test, kind="while")
|
421
|
+
pre_versions = dict(versions)
|
422
|
+
pre_latest = dict(latest)
|
423
|
+
body_ops_start = len(ops)
|
424
|
+
versions_body = dict(pre_versions)
|
425
|
+
latest_body = dict(pre_latest)
|
426
|
+
saved_versions, saved_latest = versions, latest
|
427
|
+
saved_ctx = context_suffix
|
428
|
+
ctx_counts["while"] += 1
|
429
|
+
context_suffix = f"while{ctx_counts['while']}"
|
430
|
+
versions, latest = versions_body, latest_body
|
431
|
+
for inner in stmt.body:
|
432
|
+
_parse_stmt(inner)
|
433
|
+
versions_body, latest_body = versions, latest
|
434
|
+
versions, latest = saved_versions, saved_latest
|
435
|
+
context_suffix = saved_ctx
|
436
|
+
if len(ops) > body_ops_start:
|
437
|
+
ops[body_ops_start]["deps"] = [*ops[body_ops_start].get("deps", []), cond_id]
|
438
|
+
changed = {k for k in latest_body if pre_latest.get(k) != latest_body.get(k)}
|
439
|
+
carried = [k for k in changed if k in pre_latest]
|
440
|
+
for var in sorted(carried):
|
441
|
+
phi_id = _ssa_new(var)
|
442
|
+
ops.append({
|
443
|
+
"id": phi_id,
|
444
|
+
"op": "PHI",
|
445
|
+
"deps": [pre_latest[var], latest_body[var]],
|
446
|
+
"args": {"var": var},
|
447
|
+
})
|
448
|
+
return None
|
282
449
|
elif isinstance(stmt, (ast.Pass,)):
|
283
|
-
|
450
|
+
return None
|
284
451
|
else:
|
285
|
-
raise DSLParseError("Only assignments,
|
452
|
+
raise DSLParseError("Only assignments, control flow, settings/output calls, and return are allowed in function body")
|
453
|
+
|
454
|
+
# Parse body sequentially; still require a resulting output
|
455
|
+
for i, stmt in enumerate(fn.body): # type: ignore[attr-defined]
|
456
|
+
_parse_stmt(stmt)
|
286
457
|
|
287
458
|
if not outputs:
|
288
459
|
if returned_var is not None:
|
289
460
|
outputs.append({"from": returned_var, "as": "return"})
|
290
461
|
else:
|
291
462
|
raise DSLParseError("At least one output() call required")
|
292
|
-
if len(ops) >
|
463
|
+
if len(ops) > 2000:
|
293
464
|
raise DSLParseError("Too many operations")
|
294
465
|
|
295
|
-
# Include the parsed function name for visibility/debugging
|
296
466
|
fn_name = getattr(fn, "name", None) # type: ignore[attr-defined]
|
297
|
-
plan: Dict[str, Any] = {"version":
|
467
|
+
plan: Dict[str, Any] = {"version": 2, "function": fn_name, "ops": ops, "outputs": outputs}
|
298
468
|
if settings:
|
299
469
|
plan["settings"] = settings
|
300
470
|
return plan
|
@@ -0,0 +1,11 @@
|
|
1
|
+
py2dag/__init__.py,sha256=i8VB44JCVcRJAcvnQtbH8YVRUz5j7dE355iRbikXPGQ,250
|
2
|
+
py2dag/cli.py,sha256=BCBi5mNxOqeEN8uEMt_hiDx0iSt7ZE3y74cGXREzZ2I,1296
|
3
|
+
py2dag/export_dagre.py,sha256=X2zXgHYX75-9CSEvFoOx12BbLN3X9VTFD6WYuKAL_yM,4109
|
4
|
+
py2dag/export_svg.py,sha256=YyjqOuj8GhUTDWP70SKnnSWAKI1PvJwyOhHLwB29uNM,1812
|
5
|
+
py2dag/parser.py,sha256=5_EhPzzl2P5inqT0F_JMSkPXecPzHbkQvRqz_sia_bw,22915
|
6
|
+
py2dag/pseudo.py,sha256=NJK61slyFLtSjhj8gJDJneUInEpBN57_41g8IfHNPWI,922
|
7
|
+
py2dag-0.2.0.dist-info/LICENSE,sha256=3Qee1EPwej_nusovTbyIQ8LvD2rXHdM0c6LNwk_D8Kc,1067
|
8
|
+
py2dag-0.2.0.dist-info/METADATA,sha256=8AhuHIpEfSkZsa4VhpbLRuAOWS9Gk6nx7Kt6ajRdP6Y,3549
|
9
|
+
py2dag-0.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
|
10
|
+
py2dag-0.2.0.dist-info/entry_points.txt,sha256=Q0SHexJJ0z1te4AYL1xTZogx5FrxCCE1ZJ5qntkFMZs,42
|
11
|
+
py2dag-0.2.0.dist-info/RECORD,,
|
py2dag-0.1.14.dist-info/RECORD
DELETED
@@ -1,11 +0,0 @@
|
|
1
|
-
py2dag/__init__.py,sha256=i8VB44JCVcRJAcvnQtbH8YVRUz5j7dE355iRbikXPGQ,250
|
2
|
-
py2dag/cli.py,sha256=BCBi5mNxOqeEN8uEMt_hiDx0iSt7ZE3y74cGXREzZ2I,1296
|
3
|
-
py2dag/export_dagre.py,sha256=S244wUxBMuM9qXD8bklaslkMp3rcbBGWvMhwWVBzBF0,3487
|
4
|
-
py2dag/export_svg.py,sha256=YyjqOuj8GhUTDWP70SKnnSWAKI1PvJwyOhHLwB29uNM,1812
|
5
|
-
py2dag/parser.py,sha256=y-bZk6Sb6pcOHbKdO0RJILK0V9GpBlAzsJlpEqDvoG8,16585
|
6
|
-
py2dag/pseudo.py,sha256=NJK61slyFLtSjhj8gJDJneUInEpBN57_41g8IfHNPWI,922
|
7
|
-
py2dag-0.1.14.dist-info/LICENSE,sha256=3Qee1EPwej_nusovTbyIQ8LvD2rXHdM0c6LNwk_D8Kc,1067
|
8
|
-
py2dag-0.1.14.dist-info/METADATA,sha256=ylRalH71BxB9ti6RAM5D4JbqR4kZQuOE9_oySImGr3w,3550
|
9
|
-
py2dag-0.1.14.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
|
10
|
-
py2dag-0.1.14.dist-info/entry_points.txt,sha256=Q0SHexJJ0z1te4AYL1xTZogx5FrxCCE1ZJ5qntkFMZs,42
|
11
|
-
py2dag-0.1.14.dist-info/RECORD,,
|
File without changes
|
File without changes
|
File without changes
|