py2dag 0.2.3__tar.gz → 0.3.0__tar.gz

{py2dag-0.2.3 → py2dag-0.3.0}/PKG-INFO

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

py2dag-0.3.0/py2dag/cli.py

@@ -0,0 +1,376 @@
+import argparse
+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:
+    ap = argparse.ArgumentParser(description="Convert Python function plan to DAG")
+    ap.add_argument("file", help="Python file containing the plan function")
+    ap.add_argument("--func", default=None, help="Function name to parse (auto-detect if omitted)")
+    ap.add_argument("--svg", action="store_true", help="Also export plan.svg via Graphviz (requires dot)")
+    ap.add_argument("--html", action="store_true", help="Also export plan.html via Dagre (no system deps)")
+    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(graph, f, indent=2)
+    pseudo_code = pseudo_module.generate(plan)
+    with open("plan.pseudo", "w", encoding="utf-8") as f:
+        f.write(pseudo_code)
+    if args.html:
+        from . import export_dagre
+        export_dagre.export(plan, filename="plan.html")
+    elif args.svg:
+        try:
+            export_svg.export(plan, filename="plan.svg")
+        except RuntimeError as e:
+            print(f"Warning: SVG export skipped: {e}")
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main()

py2dag-0.3.0/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-0.2.3 → py2dag-0.3.0}/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-0.3.0/py2dag/export_svg.py

@@ -0,0 +1,102 @@
+from typing import Dict, Any
+
+from .colors import color_for
+
+try:
+    from graphviz import Digraph
+except Exception:  # pragma: no cover
+    Digraph = None  # type: ignore
+
+try:  # optional: only available when graphviz is installed
+    from graphviz.backend.execute import ExecutableNotFound  # type: ignore
+except Exception:  # pragma: no cover
+    ExecutableNotFound = None  # type: ignore
+
+
+def export(plan: Dict[str, Any], filename: str = "plan.svg") -> str:
+    """Export the plan as an SVG using graphviz.
+
+    Writes a true SVG file at `filename`. If Graphviz system binaries are
+    missing, raises RuntimeError with a helpful message. This avoids leaving a
+    stray DOT file named `plan.svg` when rendering fails.
+    """
+    if Digraph is None:
+        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", []))
+    op_by_id = {op["id"]: op for op in ops}
+
+    # Nodes
+    for op in ops:
+        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:
+        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_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",
+            style="filled",
+            fillcolor=color_for("output"),
+        )
+        graph.edge(out["from"], out_id)
+
+    try:
+        # Use pipe() to obtain SVG bytes directly so we only write the
+        # destination file on successful rendering.
+        svg_bytes = graph.pipe(format="svg")
+    except Exception as e:  # pragma: no cover - depends on local system
+        if ExecutableNotFound is not None and isinstance(e, ExecutableNotFound):
+            raise RuntimeError(
+                "Graphviz 'dot' executable not found. Install Graphviz (e.g., 'brew install graphviz' on macOS) "
+                "or run without --svg."
+            ) from e
+        raise
+
+    with open(filename, "wb") as f:
+        f.write(svg_bytes)
+    return filename

{py2dag-0.2.3 → py2dag-0.3.0}/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)
@@ -288,6 +311,26 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
             })
             return ssa
 
+        def _emit_assign_from_subscript(var_name: str, node: ast.Subscript) -> str:
+            # Support name[key] where key is a JSON-serialisable literal
+            base = node.value
+            if not isinstance(base, ast.Name):
+                raise DSLParseError("Subscript base must be a variable name")
+            # Extract slice expression across Python versions
+            sl = getattr(node, 'slice', None)
+            # In Python >=3.9, slice is the actual node; before it may be ast.Index
+            if hasattr(ast, 'Index') and isinstance(sl, getattr(ast, 'Index')):  # type: ignore[attr-defined]
+                sl = sl.value  # type: ignore[assignment]
+            key = _literal(sl)  # may raise if not literal
+            ssa = _ssa_new(var_name)
+            ops.append({
+                "id": ssa,
+                "op": "GET.item",
+                "deps": [_ssa_get(base.id)],
+                "args": {"key": key},
+            })
+            return ssa
+
         def _emit_cond(node: ast.AST, kind: str = "if") -> str:
             expr = _stringify(node)
             deps = [_ssa_get(n) for n in _collect_value_deps(node)]
@@ -312,22 +355,28 @@ 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)):
                     return _emit_assign_from_literal_or_pack(var_name, value)
                 elif isinstance(value, (ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp)):
                     return _emit_assign_from_comp(var_name, value)
+                elif isinstance(value, ast.Subscript):
+                    return _emit_assign_from_subscript(var_name, value)
                 else:
                     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)
@@ -354,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):
@@ -483,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]
@@ -526,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):
@@ -538,7 +593,11 @@ 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")

{py2dag-0.2.3 → py2dag-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "py2dag"
-version = "0.2.3"
+version = "0.3.0"
 description = "Convert Python function plans to DAG (JSON, pseudo, optional SVG)."
 authors = ["rvergis"]
 license = "MIT"

py2dag-0.2.3/py2dag/cli.py

@@ -1,34 +0,0 @@
-import argparse
-import json
-
-from . import parser as dsl_parser
-from . import pseudo as pseudo_module
-from . import export_svg
-
-
-def main() -> None:
-    ap = argparse.ArgumentParser(description="Convert Python function plan to DAG")
-    ap.add_argument("file", help="Python file containing the plan function")
-    ap.add_argument("--func", default=None, help="Function name to parse (auto-detect if omitted)")
-    ap.add_argument("--svg", action="store_true", help="Also export plan.svg via Graphviz (requires dot)")
-    ap.add_argument("--html", action="store_true", help="Also export plan.html via Dagre (no system deps)")
-    args = ap.parse_args()
-
-    plan = dsl_parser.parse_file(args.file, function_name=args.func)
-    with open("plan.json", "w", encoding="utf-8") as f:
-        json.dump(plan, f, indent=2)
-    pseudo_code = pseudo_module.generate(plan)
-    with open("plan.pseudo", "w", encoding="utf-8") as f:
-        f.write(pseudo_code)
-    if args.html:
-        from . import export_dagre
-        export_dagre.export(plan, filename="plan.html")
-    elif args.svg:
-        try:
-            export_svg.export(plan, filename="plan.svg")
-        except RuntimeError as e:
-            print(f"Warning: SVG export skipped: {e}")
-
-
-if __name__ == "__main__":  # pragma: no cover
-    main()

py2dag-0.2.3/py2dag/export_svg.py

@@ -1,65 +0,0 @@
-from typing import Dict, Any
-
-try:
-    from graphviz import Digraph
-except Exception:  # pragma: no cover
-    Digraph = None  # type: ignore
-
-try:  # optional: only available when graphviz is installed
-    from graphviz.backend.execute import ExecutableNotFound  # type: ignore
-except Exception:  # pragma: no cover
-    ExecutableNotFound = None  # type: ignore
-
-
-def export(plan: Dict[str, Any], filename: str = "plan.svg") -> str:
-    """Export the plan as an SVG using graphviz.
-
-    Writes a true SVG file at `filename`. If Graphviz system binaries are
-    missing, raises RuntimeError with a helpful message. This avoids leaving a
-    stray DOT file named `plan.svg` when rendering fails.
-    """
-    if Digraph is None:
-        raise RuntimeError("Python package 'graphviz' is required for SVG export")
-
-    graph = Digraph(format="svg")
-
-    # Index ops for edge label decisions
-    ops = list(plan.get("ops", []))
-    op_by_id = {op["id"]: op for op in ops}
-
-    # Nodes
-    for op in ops:
-        graph.node(op["id"], label=op["op"])
-
-    # Dependency edges with labels showing data/control
-    for op in ops:
-        for dep in op.get("deps", []):
-            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)
-    for out in plan.get("outputs", []):
-        out_id = f"out:{out['as']}"
-        graph.node(out_id, label=out['as'], shape="note")
-        graph.edge(out["from"], out_id)
-
-    try:
-        # Use pipe() to obtain SVG bytes directly so we only write the
-        # destination file on successful rendering.
-        svg_bytes = graph.pipe(format="svg")
-    except Exception as e:  # pragma: no cover - depends on local system
-        if ExecutableNotFound is not None and isinstance(e, ExecutableNotFound):
-            raise RuntimeError(
-                "Graphviz 'dot' executable not found. Install Graphviz (e.g., 'brew install graphviz' on macOS) "
-                "or run without --svg."
-            ) from e
-        raise
-
-    with open(filename, "wb") as f:
-        f.write(svg_bytes)
-    return filename