py2dag 0.3.3__tar.gz → 0.3.5__tar.gz

{py2dag-0.3.3 → py2dag-0.3.5}/PKG-INFO

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

{py2dag-0.3.3 → py2dag-0.3.5}/py2dag/parser.py

@@ -63,6 +63,7 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
         latest: Dict[str, str] = {}
         context_suffix: str = ""
         ctx_counts: Dict[str, int] = {"if": 0, "loop": 0, "while": 0, "except": 0}
+        loop_depth = 0
 
         def _ssa_new(name: str) -> str:
             if not VALID_NAME_RE.match(name):
@@ -331,6 +332,29 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
             })
             return ssa
 
+        def _emit_assign_to_subscript(target: ast.Subscript, value: ast.AST) -> str:
+            """Emit a SET.item node for an assignment like ``name[key] = value``."""
+            base = target.value
+            if not isinstance(base, ast.Name):
+                raise DSLParseError("Subscript base must be a variable name")
+            # Extract slice expression across Python versions
+            sl = getattr(target, "slice", None)
+            if hasattr(ast, "Index") and isinstance(sl, getattr(ast, "Index")):  # type: ignore[attr-defined]
+                sl = sl.value  # type: ignore[assignment]
+            key = _literal(sl)
+
+            val_id = _emit_value(f"{base.id}_item", value)
+
+            base_id = _ssa_get(base.id)
+            ssa = _ssa_new(base.id)
+            ops.append({
+                "id": ssa,
+                "op": "SET.item",
+                "deps": [base_id, val_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)]
@@ -338,6 +362,51 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
             ops.append({"id": ssa, "op": "COND.eval", "deps": deps, "args": {"expr": expr, "kind": kind}})
             return ssa
 
+        def _emit_assign_from_ifexp(var_name: str, node: ast.IfExp) -> str:
+            """Emit operations for an inline ``a if cond else b`` expression."""
+            cond_id = _emit_cond(node.test, kind="ifexp")
+
+            then_start = len(ops)
+            then_id = _emit_value(f"{var_name}_then", node.body)
+            if len(ops) > then_start:
+                first = ops[then_start]
+                deps0 = first.get("deps", []) or []
+                if cond_id not in deps0:
+                    first["deps"] = [*deps0, cond_id]
+
+            else_start = len(ops)
+            else_id = _emit_value(f"{var_name}_else", node.orelse)
+            if len(ops) > else_start:
+                first = ops[else_start]
+                deps0 = first.get("deps", []) or []
+                if cond_id not in deps0:
+                    first["deps"] = [*deps0, cond_id]
+
+            ssa = _ssa_new(var_name)
+            ops.append({"id": ssa, "op": "PHI", "deps": [then_id, else_id], "args": {"var": var_name}})
+            return ssa
+
+        def _emit_value(var_name: str, value: ast.AST) -> str:
+            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, awaited)
+            if isinstance(value, ast.JoinedStr):
+                return _emit_assign_from_fstring(var_name, value)
+            if isinstance(value, (ast.Constant, ast.List, ast.Tuple, ast.Dict)):
+                return _emit_assign_from_literal_or_pack(var_name, value)
+            if isinstance(value, (ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp)):
+                return _emit_assign_from_comp(var_name, value)
+            if isinstance(value, ast.Subscript):
+                return _emit_assign_from_subscript(var_name, value)
+            if isinstance(value, ast.IfExp):
+                return _emit_assign_from_ifexp(var_name, value)
+            if isinstance(value, ast.Name):
+                return _ssa_get(value.id)
+            raise DSLParseError("Right hand side must be a call or f-string")
+
         def _emit_iter(node: ast.AST, target_label: Optional[str] = None) -> str:
             expr = _stringify(node)
             deps = [_ssa_get(n) for n in _collect_value_deps(node)]
@@ -349,28 +418,18 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
             return ssa
 
         def _parse_stmt(stmt: ast.stmt) -> Optional[str]:
-            nonlocal returned_var, versions, latest, context_suffix
+            nonlocal returned_var, versions, latest, context_suffix, loop_depth
             if isinstance(stmt, ast.Assign):
-                if len(stmt.targets) != 1 or not isinstance(stmt.targets[0], ast.Name):
-                    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, 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)
+                if len(stmt.targets) != 1:
+                    raise DSLParseError("Assignment must have exactly one target")
+                target = stmt.targets[0]
+                if isinstance(target, ast.Name):
+                    var_name = target.id
+                    return _emit_value(var_name, stmt.value)
+                elif isinstance(target, ast.Subscript):
+                    return _emit_assign_to_subscript(target, stmt.value)
                 else:
-                    raise DSLParseError("Right hand side must be a call or f-string")
+                    raise DSLParseError("Assignment targets must be simple names")
             elif isinstance(stmt, ast.Expr):
                 call = stmt.value
                 awaited = False
@@ -406,8 +465,22 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                     _emit_expr_call(call, awaited)
                 return None
             elif isinstance(stmt, ast.Return):
+                if loop_depth > 0:
+                    ssa = _ssa_new("break")
+                    ops.append({"id": ssa, "op": "CTRL.break", "deps": [], "args": {}})
                 if isinstance(stmt.value, ast.Name):
-                    returned_var = _ssa_get(stmt.value.id)
+                    name = stmt.value.id
+                    if name in latest:
+                        returned_var = _ssa_get(name)
+                    else:
+                        const_id = _ssa_new("return_value")
+                        ops.append({
+                            "id": const_id,
+                            "op": "CONST.value",
+                            "deps": [],
+                            "args": {"value": None},
+                        })
+                        returned_var = const_id
                 elif isinstance(stmt.value, (ast.Constant, ast.List, ast.Tuple, ast.Dict)):
                     lit = _literal(stmt.value)
                     const_id = _ssa_new("return_value")
@@ -503,6 +576,7 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                 saved_ctx = context_suffix
                 ctx_counts["loop"] += 1
                 context_suffix = f"loop{ctx_counts['loop']}"
+                loop_depth += 1
                 versions, latest = versions_body, latest_body
                 # Predefine loop target variables as items from iterator for dependency resolution
                 def _bind_loop_target(target: ast.AST):
@@ -529,6 +603,7 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                 _bind_loop_target(stmt.target)
                 for inner in stmt.body:
                     _parse_stmt(inner)
+                loop_depth -= 1
                 versions_body, latest_body = versions, latest
                 versions, latest = saved_versions, saved_latest
                 context_suffix = saved_ctx
@@ -571,9 +646,11 @@ def parse(source: str, function_name: Optional[str] = None) -> Dict[str, Any]:
                 saved_ctx = context_suffix
                 ctx_counts["while"] += 1
                 context_suffix = f"while{ctx_counts['while']}"
+                loop_depth += 1
                 versions, latest = versions_body, latest_body
                 for inner in stmt.body:
                     _parse_stmt(inner)
+                loop_depth -= 1
                 versions_body, latest_body = versions, latest
                 versions, latest = saved_versions, saved_latest
                 context_suffix = saved_ctx

{py2dag-0.3.3 → py2dag-0.3.5}/pyproject.toml

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