py2dag 0.2.4__py3-none-any.whl → 0.3.1__py3-none-any.whl

py2dag/cli.py

@@ -4,6 +4,346 @@ import json
 from . import parser as dsl_parser
 from . import pseudo as pseudo_module
 from . import export_svg
+import json as _json
+import re as _re
+
+
+def _to_nodes_edges(plan: dict) -> dict:
+    """Convert internal plan (with ops/outputs) to explicit nodes/edges.
+
+    Nodes: one per op, plus one per output sink.
+    Edges: for each dep -> op.id, with labels similar to exporters; and op.id -> output sink.
+    """
+    ops = list(plan.get("ops", []))
+    outputs = list(plan.get("outputs", []))
+
+    op_by_id = {op["id"]: op for op in ops}
+
+    # Optional remapping of node ids for nicer presentation
+    base_id_map: dict[str, str] = {}
+    loop_ctx_map: dict[str, str] = {}  # maps 'loop1' -> 'for_loop_1'
+    for_count = 0
+    for op in ops:
+        new_id = op["id"]
+        if op.get("op") == "ITER.eval" and (op.get("args") or {}).get("kind") == "for":
+            for_count += 1
+            new_id = f"for_loop_{for_count}"
+            loop_ctx_map[f"loop{for_count}"] = new_id
+        base_id_map[op["id"]] = new_id
+
+    def _remap_ctx_suffix(op_id: str) -> str:
+        # Replace '@loopN' with '@for_loop_N' when present
+        if '@' in op_id:
+            base, ctx = op_id.split('@', 1)
+            mapped_ctx = loop_ctx_map.get(ctx)
+            if mapped_ctx:
+                return f"{base}@{mapped_ctx}"
+        return op_id
+
+    # Final id map includes context-aware remapping
+    id_map: dict[str, str] = {}
+    for op in ops:
+        oid = op["id"]
+        # first apply context suffix remap, then base remap if any
+        ctx_remapped = _remap_ctx_suffix(oid)
+        id_map[oid] = base_id_map.get(ctx_remapped, ctx_remapped)
+
+    nodes = []
+    edges = []
+
+    def _base_name(ssa: str) -> str:
+        # Extract base variable name from SSA id like "name_2@ctx" -> "name"
+        m = _re.match(r"^([a-z_][a-z0-9_]*?)_\d+(?:@.*)?$", ssa)
+        return m.group(1) if m else ssa
+
+    def _expr_for(op: dict) -> str:
+        op_name = op.get("op", "")
+        deps = op.get("deps", []) or []
+        args = op.get("args", {}) or {}
+        dep_labels = op.get("dep_labels", []) or []
+        # Helper to stringify deps as base names
+        dep_names = [_base_name(d) for d in deps]
+        # Special-case control/struct ops first
+        if op_name == "TEXT.format":
+            return f"f\"{args.get('template','')}\""
+        if op_name == "CONST.value":
+            return _json.dumps(args.get("value"))
+        if op_name == "GET.item":
+            base = _base_name(deps[0]) if deps else "<unknown>"
+            return f"{base}[{_json.dumps(args.get('key'))}]"
+        if op_name == "PACK.list":
+            return "[" + ", ".join(dep_names) + "]"
+        if op_name == "PACK.tuple":
+            return "(" + ", ".join(dep_names) + ")"
+        if op_name == "PACK.dict":
+            keys = (args.get("keys") or [])
+            items = []
+            for i, k in enumerate(keys):
+                val = dep_names[i] if i < len(dep_names) else "<val>"
+                items.append(f"{k}: {val}")
+            return "{" + ", ".join(items) + "}"
+        if op_name == "COND.eval":
+            kind = args.get("kind")
+            return f"{kind} {args.get('expr')}" if kind else str(args.get('expr'))
+        if op_name == "ITER.eval":
+            tgt = args.get("target")
+            expr = args.get("expr")
+            if tgt:
+                return f"for {tgt} in {expr}"
+            return f"iter {expr}"
+        if op_name == "ITER.item":
+            return str(args.get("target") or "item")
+        if op_name == "PHI":
+            return "phi(" + ", ".join(dep_names) + ")"
+        if op_name.startswith("COMP."):
+            return op_name
+        # Generic tool/attribute call ops
+        if op_name.endswith(".op") or "." in op_name:
+            parts = []
+            used_args = set()
+            # Encode deps in original order, honoring labels for kwargs and vararg markers
+            for i, name in enumerate(dep_names):
+                lbl = dep_labels[i] if i < len(dep_labels) else ""
+                dep_id = deps[i] if i < len(deps) else ""
+                src = op_by_id.get(dep_id)
+                if (
+                    lbl
+                    and src
+                    and src.get("op") == "CONST.value"
+                    and isinstance(args, dict)
+                    and lbl in args
+                ):
+                    parts.append(f"{lbl}={_json.dumps(args[lbl])}")
+                    used_args.add(lbl)
+                elif lbl == "*":
+                    parts.append(f"*{name}")
+                elif lbl == "**":
+                    parts.append(f"**{name}")
+                elif lbl:  # keyword variable arg
+                    parts.append(f"{lbl}={name}")
+                else:  # positional
+                    parts.append(name)
+            # Append literal kwargs (in insertion order)
+            for k, v in (args.items() if isinstance(args, dict) else []):
+                if k in {"template", "expr", "kind", "keys", "target", "var"} or k in used_args:
+                    continue
+                parts.append(f"{k}={_json.dumps(v)}")
+            return f"{op_name}(" + ", ".join(parts) + ")"
+        # Fallback: show op name and dep names
+        return f"{op_name}(" + ", ".join(dep_names) + ")"
+
+    def _type_for(op: dict) -> str:
+        op_name = op.get("op", "")
+        if op_name == "CONST.value":
+            return "const"
+        if op_name == "TEXT.format":
+            return "format"
+        if op_name == "GET.item":
+            return "get_item"
+        if op_name.startswith("PACK."):
+            kind = op_name.split(".", 1)[1]
+            return f"pack:{kind}"
+        if op_name == "COND.eval":
+            kind = (op.get("args") or {}).get("kind")
+            if kind in {"if", "while"}:
+                return kind
+            return "cond"
+        if op_name == "ITER.eval":
+            kind = (op.get("args") or {}).get("kind")
+            return "forloop" if kind == "for" else "iter"
+        if op_name == "ITER.item":
+            return "forloop_item"
+        if op_name == "PHI":
+            return "phi"
+        if op_name == "CTRL.break":
+            return "break"
+        if op_name.startswith("COMP."):
+            comp = op_name.split(".", 1)[1]
+            return f"comp:{comp}"
+        # Default: treat as a call/tool node
+        return "call"
+
+    # Create nodes for ops
+    # For friendly IDs, keep execution-order counters per base
+    pretty_counters: dict[str, int] = {}
+
+    for op in ops:
+        # Compute final node id (may override for pretty IDs)
+        op_id = op["id"]
+        op_name = op.get("op", "")
+        # Choose displayed variable name: prefer loop target for ITER.eval
+        var_name = _base_name(op_id)
+        if op.get("op") == "ITER.eval":
+            tgt = (op.get("args") or {}).get("target")
+            if isinstance(tgt, str) and tgt:
+                var_name = tgt
+        # Start from base remapped id
+        node_id = id_map[op_id]
+        # If this is a synthetic expression-call (base var 'call'), use op name plus context suffix
+        if var_name == "call" and (op_name.endswith(".op") or "." in op_name):
+            suffix = ""
+            if "@" in node_id:
+                suffix = node_id.split("@", 1)[1]
+            base_pretty = f"{op_name}@{suffix}" if suffix else op_name
+            # Increment execution-order counter and append #N, starting at 1
+            n = pretty_counters.get(base_pretty, 0) + 1
+            pretty_counters[base_pretty] = n
+            node_id = f"{base_pretty}#{n}"
+            # Update map so downstream references use the pretty id
+            id_map[op_id] = node_id
+
+        # Merge literal kwargs with variable bindings into args map
+        merged_args = dict(op.get("args", {}) or {})
+        dep_labels = op.get("dep_labels", []) or []
+        # track positional index for unlabeled deps
+        pos_index = 0
+        for idx, dep in enumerate(op.get("deps", []) or []):
+            label = dep_labels[idx] if idx < len(dep_labels) else ""
+            if not label:
+                # positional arg
+                merged_args[str(pos_index)] = dep
+                pos_index += 1
+            elif label not in {"*", "**"}:
+                # keyword variable arg
+                merged_args[label] = dep
+            # ignore '*'/'**' labels in args map; they are reflected by numeric positions or literal kwargs
+
+        # Remap any arg values that reference op ids
+        for k, v in list(merged_args.items()):
+            if isinstance(v, str):
+                vv = id_map.get(v, _remap_ctx_suffix(v))
+                merged_args[k] = vv
+
+        # Ensure deterministic ordering of params in args (sorted by key)
+        if merged_args:
+            merged_args = {k: merged_args[k] for k in sorted(merged_args)}
+
+        # Choose displayed variable name: prefer loop target for ITER.eval
+        var_name = _base_name(op["id"])
+        if op.get("op") == "ITER.eval":
+            tgt = (op.get("args") or {}).get("target")
+            if isinstance(tgt, str) and tgt:
+                var_name = tgt
+
+        label = op_name
+        if op.get("op") == "ITER.eval" and (op.get("args") or {}).get("kind") == "for":
+            label = "forloop"
+        elif op.get("op") == "ITER.item":
+            label = "FORLOOP.item"
+
+        node_obj = {
+            "id": node_id,
+            "type": _type_for(op),
+            "label": label,
+            "var": var_name,
+            "expr": _expr_for(op),
+            "args": merged_args,
+        }
+        # For control/struct item binders, drop args to avoid redundant control deps
+        if op.get("op") == "ITER.item":
+            node_obj.pop("args", None)
+        nodes.append(node_obj)
+
+    # Create edges for deps
+    for op in ops:
+        deps = op.get("deps", []) or []
+        # Special case: PACK.dict can label edges by key
+        keys = []
+        if op.get("op") == "PACK.dict":
+            keys = list((op.get("args") or {}).get("keys", []) or [])
+        dep_labels = op.get("dep_labels", []) or []
+        seen = set()
+        for idx, dep in enumerate(deps):
+            to_id = id_map[op["id"]]
+            from_id = id_map.get(dep, _remap_ctx_suffix(dep))
+            if (from_id, to_id) in seen:
+                continue
+            seen.add((from_id, to_id))
+            label = (dep_labels[idx] if idx < len(dep_labels) else "") or ""
+            src = op_by_id.get(dep)
+            if not label:
+                if keys and idx < len(keys):
+                    label = str(keys[idx])
+                elif src is not None:
+                    if src.get("op") == "COND.eval":
+                        # Label IF branches as then/else based on destination context; keep 'cond' for while
+                        kind = (src.get("args") or {}).get("kind")
+                        if kind == "if":
+                            if "@then" in to_id:
+                                label = "then"
+                            elif "@else" in to_id:
+                                label = "else"
+                            else:
+                                label = "cond"
+                        else:
+                            label = "cond"
+                    elif src.get("op") == "ITER.eval":
+                        label = str((src.get("args") or {}).get("target") or "iter")
+            edges.append({"from": from_id, "to": to_id, "label": label})
+
+    # Add explicit true/false edges for `if ...: continue` patterns inside loops
+    # When an IF inside a loop has a `continue` body and no explicit else branch,
+    # no edges are emitted from the COND node to represent the two control-flow
+    # outcomes.  Detect such cases and add synthetic edges so that both the
+    # "true" (continue) and "false" paths are visible in the exported graph.
+    for idx, op in enumerate(ops):
+        if op.get("op") != "COND.eval":
+            continue
+        oid = op["id"]
+        if "@" not in oid:
+            continue
+        _, ctx = oid.split("@", 1)
+        if not ctx.startswith("loop"):
+            continue
+        from_id = id_map[oid]
+        existing = [e for e in edges if e["from"] == from_id]
+        # Skip if we already have explicit branch labels from this cond
+        if any(e.get("label") in {"then", "else", "true", "false"} for e in existing):
+            continue
+        # Find first op after this cond within the same loop context
+        next_op = None
+        for later in ops[idx + 1 :]:
+            if later["id"].endswith(f"@{ctx}"):
+                next_op = later
+                break
+        if next_op is not None:
+            to_id = id_map[next_op["id"]]
+            # If an unlabeled edge already exists, relabel it as the false branch
+            updated = False
+            for e in existing:
+                if e["to"] == to_id:
+                    e["label"] = "false"
+                    updated = True
+                    break
+            if not updated:
+                edges.append({"from": from_id, "to": to_id, "label": "false"})
+        # True branch loops back to the for-loop node itself
+        loop_node = loop_ctx_map.get(ctx)
+        if loop_node:
+            edges.append({"from": from_id, "to": loop_node, "label": "true"})
+
+    # Output nodes and edges
+    for out in outputs:
+        out_id = f"out:{out['as']}"
+        nodes.append({
+            "id": out_id,
+            "type": "output",
+            "label": out["as"],
+            "var": out["as"],
+            "expr": out["as"],
+            "args": {"kind": "output"},
+        })
+        edges.append({"from": out["from"], "to": out_id})
+
+    graph = {
+        "version": plan.get("version", 2),
+        "function": plan.get("function"),
+        "nodes": nodes,
+        "edges": edges,
+    }
+    if plan.get("settings"):
+        graph["settings"] = plan["settings"]
+    return graph
 
 
 def main() -> None:
@@ -15,8 +355,10 @@ def main() -> None:
     args = ap.parse_args()
 
     plan = dsl_parser.parse_file(args.file, function_name=args.func)
+    # Write explicit nodes/edges form to plan.json for downstream use
+    graph = _to_nodes_edges(plan)
     with open("plan.json", "w", encoding="utf-8") as f:
-        json.dump(plan, f, indent=2)
+        json.dump(graph, f, indent=2)
     pseudo_code = pseudo_module.generate(plan)
     with open("plan.pseudo", "w", encoding="utf-8") as f:
         f.write(pseudo_code)

py2dag/colors.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import hashlib
+
+CRAYON_COLORS = [
+    "cornflowerblue",
+    "lightcoral",
+    "gold",
+    "mediumseagreen",
+    "orchid",
+    "sandybrown",
+    "plum",
+    "turquoise",
+    "khaki",
+    "salmon",
+]
+
+
+def color_for(name: str) -> str:
+    """Return a pseudo-random but stable color for a given name.
+
+    Picks from CRAYON_COLORS using a SHA256 hash for determinism so the same
+    node type is colored consistently across exports.
+    """
+    h = hashlib.sha256(name.encode("utf-8")).hexdigest()
+    idx = int(h, 16) % len(CRAYON_COLORS)
+    return CRAYON_COLORS[idx]
+

py2dag/export_dagre.py

@@ -4,6 +4,8 @@ import json
 from pathlib import Path
 from typing import Any, Dict
 
+from .colors import color_for
+
 
 HTML_TEMPLATE = """<!doctype html>
 <html lang="en">
@@ -15,7 +17,7 @@ HTML_TEMPLATE = """<!doctype html>
     body { font-family: -apple-system, BlinkMacSystemFont, Segoe UI, Roboto, Arial, sans-serif; margin: 0; padding: 0; }
     header { padding: 10px 16px; background: #111; color: #eee; font-size: 14px; }
     #container { padding: 12px; }
-    svg { width: 100%; height: 80vh; border-top: 1px solid #ddd; }
+    svg { width: 100%; height: 80vh; border: 1px solid #ddd; margin: 10px; padding: 10px; }
     .node rect { stroke: #666; fill: #fff; rx: 4; ry: 4; }
     .node.note rect { fill: #fff8dc; }
     .edgePath path { stroke: #333; fill: none; stroke-width: 1.2px; }
@@ -30,6 +32,7 @@ HTML_TEMPLATE = """<!doctype html>
   </div>
   <script>
     const plan = __PLAN_JSON__;
+    const COLOR_MAP = __COLOR_MAP__;
 
     function showMessage(msg) {
       const el = document.getElementById('container');
@@ -44,7 +47,7 @@ HTML_TEMPLATE = """<!doctype html>
     } else {
       try {
         const g = new dagreD3.graphlib.Graph({ multigraph: true })
-          .setGraph({ rankdir: 'LR', nodesep: 30, ranksep: 40 });
+          .setGraph({ rankdir: 'TB', nodesep: 30, ranksep: 40 });
         // Ensure edges have an object for labels/attrs to avoid TypeErrors
         g.setDefaultEdgeLabel(() => ({}));
 
@@ -63,13 +66,15 @@ HTML_TEMPLATE = """<!doctype html>
             label = 'PHI' + (op.args && op.args.var ? ` (${op.args.var})` : '');
             klass = 'note';
           }
-          g.setNode(op.id, { label, class: klass, padding: 8 });
+          const color = COLOR_MAP[op.op] || '#fff';
+          g.setNode(op.id, { label, class: klass, padding: 8, style: 'fill: ' + color });
         });
 
         // Add output nodes and edges from source to output
         (plan.outputs || []).forEach(out => {
           const outId = `out:${out.as}`;
-          g.setNode(outId, { label: out.as, class: 'note', padding: 8 });
+          const ocolor = COLOR_MAP['output'] || '#fff';
+          g.setNode(outId, { label: out.as, class: 'note', padding: 8, style: 'fill: ' + ocolor });
           g.setEdge(out.from, outId);
         });
 
@@ -79,13 +84,32 @@ HTML_TEMPLATE = """<!doctype html>
 
         // Add dependency edges between ops with labels
         (plan.ops || []).forEach(op => {
-          (op.deps || []).forEach(dep => {
+          const depLabels = op.dep_labels || [];
+          const seen = new Set();
+          (op.deps || []).forEach((dep, idx) => {
+            const pair = dep + '->' + op.id;
+            if (seen.has(pair)) return; seen.add(pair);
             const src = opById[dep];
-            let edgeLabel = dep; // default to SSA id
-            if (src && src.op === 'COND.eval') {
-              edgeLabel = 'cond';
-            } else if (src && src.op === 'ITER.eval') {
-              edgeLabel = (src.args && src.args.target) ? src.args.target : 'iter';
+            let edgeLabel = (depLabels[idx] || '').toString();
+            // Special-cases override when no explicit label was provided
+            if (!edgeLabel) {
+              if (op.op === 'PACK.dict' && op.args && Array.isArray(op.args.keys)) {
+                edgeLabel = (op.args.keys[idx] || '').toString();
+              } else if (src && src.op === 'COND.eval') {
+                // For IF, label branches as then/else based on dest id; for WHILE keep 'cond'
+                const kind = src.args && src.args.kind;
+                if (kind === 'if') {
+                  if ((op.id || '').includes('@then')) edgeLabel = 'then';
+                  else if ((op.id || '').includes('@else')) edgeLabel = 'else';
+                  else edgeLabel = 'cond';
+                } else {
+                  edgeLabel = 'cond';
+                }
+              } else if (src && src.op === 'ITER.eval') {
+                edgeLabel = (src.args && src.args.target) ? src.args.target : 'iter';
+              } else {
+                edgeLabel = dep; // fallback: SSA id
+              }
             }
             g.setEdge(dep, op.id, { label: edgeLabel });
           });
@@ -117,7 +141,10 @@ def export(plan: Dict[str, Any], filename: str = "plan.html") -> str:
 
     Returns the written filename.
     """
+    color_map = {op["op"]: color_for(op["op"]) for op in plan.get("ops", [])}
+    color_map["output"] = color_for("output")
     html = HTML_TEMPLATE.replace("__PLAN_JSON__", json.dumps(plan))
+    html = html.replace("__COLOR_MAP__", json.dumps(color_map))
     path = Path(filename)
     path.write_text(html, encoding="utf-8")
     return str(path)

py2dag/export_svg.py

@@ -1,5 +1,7 @@
 from typing import Dict, Any
 
+from .colors import color_for
+
 try:
     from graphviz import Digraph
 except Exception:  # pragma: no cover
@@ -22,6 +24,9 @@ def export(plan: Dict[str, Any], filename: str = "plan.svg") -> str:
         raise RuntimeError("Python package 'graphviz' is required for SVG export")
 
     graph = Digraph(format="svg")
+    # Top-down layout with a bounding box around the graph
+    graph.attr(rankdir="TB")
+    graph.attr("graph", margin="0.2", pad="0.3", color="#bbb")
 
     # Index ops for edge label decisions
     ops = list(plan.get("ops", []))
@@ -29,23 +34,55 @@ def export(plan: Dict[str, Any], filename: str = "plan.svg") -> str:
 
     # Nodes
     for op in ops:
-        graph.node(op["id"], label=op["op"])
+        graph.node(
+            op["id"],
+            label=op["op"],
+            style="filled",
+            fillcolor=color_for(op["op"]),
+        )
 
     # Dependency edges with labels showing data/control
     for op in ops:
-        for dep in op.get("deps", []):
+        dep_labels = op.get("dep_labels", []) or []
+        seen = set()
+        for idx, dep in enumerate(op.get("deps", []) or []):
+            pair = (dep, op["id"]) 
+            if pair in seen:
+                continue
+            seen.add(pair)
             src = op_by_id.get(dep)
-            label = dep  # default to SSA id
-            if src is not None:
-                if src.get("op") == "COND.eval":
-                    label = "cond"
-                elif src.get("op") == "ITER.eval":
-                    args = src.get("args", {}) or {}
-                    label = str(args.get("target") or "iter")
-            graph.edge(dep, op["id"], label=label)
+            label = (dep_labels[idx] if idx < len(dep_labels) else "") or ""
+            if not label:
+                if op.get("op") == "PACK.dict":
+                    keys = (op.get("args", {}) or {}).get("keys", []) or []
+                    if idx < len(keys):
+                        label = str(keys[idx])
+                elif src is not None:
+                    if src.get("op") == "COND.eval":
+                        const_args = src.get("args", {}) or {}
+                        if const_args.get("kind") == "if":
+                            dest_id = op.get("id", "")
+                            if "@then" in dest_id:
+                                label = "then"
+                            elif "@else" in dest_id:
+                                label = "else"
+                            else:
+                                label = "cond"
+                        else:
+                            label = "cond"
+                    elif src.get("op") == "ITER.eval":
+                        args = src.get("args", {}) or {}
+                        label = str(args.get("target") or "iter")
+            graph.edge(dep, op["id"], label=label or dep)
     for out in plan.get("outputs", []):
         out_id = f"out:{out['as']}"
-        graph.node(out_id, label=out['as'], shape="note")
+        graph.node(
+            out_id,
+            label=out['as'],
+            shape="note",
+            style="filled",
+            fillcolor=color_for("output"),
+        )
         graph.edge(out["from"], out_id)
 
     try:

py2dag/parser.py

@@ -118,9 +118,10 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
             except Exception:
                 return node.__class__.__name__
 
-        def _emit_assign_from_call(var_name: str, call: ast.Call) -> str:
+        def _emit_assign_from_call(var_name: str, call: ast.Call, awaited: bool = False) -> str:
             op_name = _get_call_name(call.func)
             deps: List[str] = []
+            dep_labels: List[str] = []
 
             def _expand_star_name(ssa_var: str) -> List[str]:
                 # Expand if previous op was a PACK.*
@@ -133,21 +134,26 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                 if isinstance(arg, ast.Starred):
                     star_val = arg.value
                     if isinstance(star_val, ast.Name):
-                        deps.extend(_expand_star_name(_ssa_get(star_val.id)))
+                        expanded = _expand_star_name(_ssa_get(star_val.id))
+                        deps.extend(expanded)
+                        dep_labels.extend(["*"] * len(expanded))
                     elif isinstance(star_val, (ast.List, ast.Tuple)):
                         for elt in star_val.elts:
                             if not isinstance(elt, ast.Name):
                                 raise DSLParseError("Starred list/tuple elements must be names")
                             deps.append(_ssa_get(elt.id))
+                            dep_labels.append("*")
                     else:
                         raise DSLParseError("*args must be a name or list/tuple of names")
                 elif isinstance(arg, ast.Name):
                     deps.append(_ssa_get(arg.id))
+                    dep_labels.append("")
                 elif isinstance(arg, (ast.List, ast.Tuple)):
                     for elt in arg.elts:
                         if not isinstance(elt, ast.Name):
                             raise DSLParseError("List/Tuple positional args must be variable names")
                         deps.append(_ssa_get(elt.id))
+                        dep_labels.append("")
                 else:
                     raise DSLParseError("Positional args must be variable names or lists/tuples of names")
 
@@ -161,21 +167,38 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                             kwargs[str(k)] = val
                     elif isinstance(v, ast.Name):
                         deps.append(_ssa_get(v.id))
+                        dep_labels.append("**")
                     else:
                         raise DSLParseError("**kwargs must be a dict literal or a variable name")
                 else:
                     if isinstance(kw.value, ast.Name):
                         deps.append(_ssa_get(kw.value.id))
+                        dep_labels.append(kw.arg or "")
                     else:
-                        kwargs[kw.arg] = _literal(kw.value)
+                        lit = _literal(kw.value)
+                        kwargs[kw.arg] = lit
+                        const_id = _ssa_new(f"{var_name}_{kw.arg}")
+                        ops.append(
+                            {
+                                "id": const_id,
+                                "op": "CONST.value",
+                                "deps": [],
+                                "args": {"value": lit},
+                            }
+                        )
+                        deps.append(const_id)
+                        dep_labels.append(kw.arg or "")
 
             ssa = _ssa_new(var_name)
-            ops.append({"id": ssa, "op": op_name, "deps": deps, "args": kwargs})
+            op: Dict[str, Any] = {"id": ssa, "op": op_name, "deps": deps, "args": kwargs, "dep_labels": dep_labels}
+            if awaited:
+                op["await"] = True
+            ops.append(op)
             return ssa
 
-        def _emit_expr_call(call: ast.Call) -> str:
+        def _emit_expr_call(call: ast.Call, awaited: bool = False) -> str:
             """Emit a node for a bare expression call (no assignment)."""
-            return _emit_assign_from_call("call", call)
+            return _emit_assign_from_call("call", call, awaited)
 
         def _emit_assign_from_fstring(var_name: str, fstr: ast.JoinedStr) -> str:
             deps: List[str] = []
@@ -241,7 +264,7 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                             inner = v_node.value
                             if not isinstance(inner, ast.Call):
                                 raise DSLParseError("await must wrap a call in dict value")
-                            tmp_id = _emit_assign_from_call(f"{var_name}_field", inner)
+                            tmp_id = _emit_assign_from_call(f"{var_name}_field", inner, awaited=True)
                             deps.append(tmp_id)
                         elif isinstance(v_node, ast.Call):
                             tmp_id = _emit_assign_from_call(f"{var_name}_field", v_node)
@@ -332,10 +355,12 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                     raise DSLParseError("Assignment targets must be simple names")
                 var_name = stmt.targets[0].id
                 value = stmt.value
+                awaited = False
                 if isinstance(value, ast.Await):
                     value = value.value
+                    awaited = True
                 if isinstance(value, ast.Call):
-                    return _emit_assign_from_call(var_name, value)
+                    return _emit_assign_from_call(var_name, value, awaited)
                 elif isinstance(value, ast.JoinedStr):
                     return _emit_assign_from_fstring(var_name, value)
                 elif isinstance(value, (ast.Constant, ast.List, ast.Tuple, ast.Dict)):
@@ -348,8 +373,10 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                     raise DSLParseError("Right hand side must be a call or f-string")
             elif isinstance(stmt, ast.Expr):
                 call = stmt.value
+                awaited = False
                 if isinstance(call, ast.Await):
                     call = call.value
+                    awaited = True
                 if not isinstance(call, ast.Call):
                     raise DSLParseError("Only call expressions allowed at top level")
                 name = _get_call_name(call.func)
@@ -376,7 +403,7 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                     outputs.append({"from": ssa_from, "as": filename})
                 else:
                     # General expression call: represent as an op node too
-                    _emit_expr_call(call)
+                    _emit_expr_call(call, awaited)
                 return None
             elif isinstance(stmt, ast.Return):
                 if isinstance(stmt.value, ast.Name):
@@ -505,9 +532,12 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                 versions_body, latest_body = versions, latest
                 versions, latest = saved_versions, saved_latest
                 context_suffix = saved_ctx
-                # Add iter dep to first op in body
+                # Add iter dep to first op in body if not already present
                 if len(ops) > body_ops_start:
-                    ops[body_ops_start]["deps"] = [*ops[body_ops_start].get("deps", []), iter_id]
+                    first = ops[body_ops_start]
+                    deps0 = first.get("deps", []) or []
+                    if iter_id not in deps0:
+                        first["deps"] = [*deps0, iter_id]
                 # Emit a summary foreach comp node depending on iterable value deps
                 iter_name_deps = _collect_value_deps(stmt.iter)
                 foreach_deps = [_ssa_get(n) for n in iter_name_deps]
@@ -548,7 +578,10 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                 versions, latest = saved_versions, saved_latest
                 context_suffix = saved_ctx
                 if len(ops) > body_ops_start:
-                    ops[body_ops_start]["deps"] = [*ops[body_ops_start].get("deps", []), cond_id]
+                    first = ops[body_ops_start]
+                    deps0 = first.get("deps", []) or []
+                    if cond_id not in deps0:
+                        first["deps"] = [*deps0, cond_id]
                 changed = {k for k in latest_body if pre_latest.get(k) != latest_body.get(k)}
                 carried = [k for k in changed if k in pre_latest]
                 for var in sorted(carried):
@@ -560,13 +593,20 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                         "args": {"var": var},
                     })
                 return None
-            elif isinstance(stmt, (ast.Pass, ast.Continue, ast.Break)):
+            elif isinstance(stmt, ast.Break):
+                ssa = _ssa_new("break")
+                ops.append({"id": ssa, "op": "CTRL.break", "deps": [], "args": {}})
+                return None
+            elif isinstance(stmt, (ast.Pass, ast.Continue)):
                 return None
             else:
                 raise DSLParseError("Only assignments, control flow, settings/output calls, and return are allowed in function body")
 
         # Parse body sequentially; still require a resulting output
-        for i, stmt in enumerate(fn.body):  # type: ignore[attr-defined]
+        body = list(getattr(fn, "body", []))  # type: ignore[attr-defined]
+        if body and isinstance(body[0], ast.Expr) and isinstance(getattr(body[0], "value", None), ast.Constant) and isinstance(getattr(body[0].value, "value", None), str):  # type: ignore[attr-defined]
+            body = body[1:]
+        for stmt in body:
             _parse_stmt(stmt)
 
         if not outputs:

{py2dag-0.2.4.dist-info → py2dag-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: py2dag
-Version: 0.2.4
+Version: 0.3.1
 Summary: Convert Python function plans to DAG (JSON, pseudo, optional SVG).
 License: MIT
 Author: rvergis

py2dag-0.3.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+py2dag/__init__.py,sha256=i8VB44JCVcRJAcvnQtbH8YVRUz5j7dE355iRbikXPGQ,250
+py2dag/cli.py,sha256=q8ocafpGtkrls1mx9QiFU1wL6KKeeJhT0Oklk6SwHB0,15151
+py2dag/colors.py,sha256=kVMVWUKJY1uLQFQux5XKiEcofDFpNPQoAEzpMJWs2b0,604
+py2dag/export_dagre.py,sha256=Og-oVjKzc7bRgeU7rgLHDjiY6wIcNyr2nIYlSAQzKoM,6041
+py2dag/export_svg.py,sha256=3ZmZqxIUidSbuKIIh5kFqhfVnoqZ6PbEQpMih4AX9xo,3793
+py2dag/parser.py,sha256=bpwOZVsJUS3fRQTFfm2CfTxqlTWUYaLYDOHJRA2Yj34,30657
+py2dag/pseudo.py,sha256=NJK61slyFLtSjhj8gJDJneUInEpBN57_41g8IfHNPWI,922
+py2dag-0.3.1.dist-info/LICENSE,sha256=3Qee1EPwej_nusovTbyIQ8LvD2rXHdM0c6LNwk_D8Kc,1067
+py2dag-0.3.1.dist-info/METADATA,sha256=GhuNW_nWmz0P0eNjX6Zs2y65AglkhUT1lplPeDzxmaU,3549
+py2dag-0.3.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+py2dag-0.3.1.dist-info/entry_points.txt,sha256=Q0SHexJJ0z1te4AYL1xTZogx5FrxCCE1ZJ5qntkFMZs,42
+py2dag-0.3.1.dist-info/RECORD,,

py2dag-0.2.4.dist-info/RECORD

@@ -1,11 +0,0 @@
-py2dag/__init__.py,sha256=i8VB44JCVcRJAcvnQtbH8YVRUz5j7dE355iRbikXPGQ,250
-py2dag/cli.py,sha256=BCBi5mNxOqeEN8uEMt_hiDx0iSt7ZE3y74cGXREzZ2I,1296
-py2dag/export_dagre.py,sha256=lScDQCIh5nE3RI_ac4Ye8AsLnn3MIn-fm1WlqP9Ei4U,4608
-py2dag/export_svg.py,sha256=tKqh16hCrCE1QJ9-MTue1ELRxmXOb2akokpSz_a5eE8,2383
-py2dag/parser.py,sha256=mJX_rjjSnO7G0-3thh5bBpPlG8N0IWAUqnp4Yn7Vd3A,28711
-py2dag/pseudo.py,sha256=NJK61slyFLtSjhj8gJDJneUInEpBN57_41g8IfHNPWI,922
-py2dag-0.2.4.dist-info/LICENSE,sha256=3Qee1EPwej_nusovTbyIQ8LvD2rXHdM0c6LNwk_D8Kc,1067
-py2dag-0.2.4.dist-info/METADATA,sha256=yNQ2eR_WmySOQGRyq0uzvuYeKzV5_AxH2FTlSfPIvd0,3549
-py2dag-0.2.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-py2dag-0.2.4.dist-info/entry_points.txt,sha256=Q0SHexJJ0z1te4AYL1xTZogx5FrxCCE1ZJ5qntkFMZs,42
-py2dag-0.2.4.dist-info/RECORD,,