pyoco 0.3.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

pyoco/dsl/expressions.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import ast
+import re
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Union
+
+
+class ExpressionSyntaxError(ValueError):
+    pass
+
+
+class ExpressionEvaluationError(RuntimeError):
+    pass
+
+
+DOT_PATH_RE = re.compile(r"^[A-Za-z_][\w.]*$")
+
+
+@dataclass(frozen=True)
+class Expression:
+    source: str
+    _python: str = field(init=False, repr=False)
+    _code: object = field(init=False, repr=False)
+
+    def __post_init__(self):
+        if not isinstance(self.source, str):
+            raise TypeError("Expression source must be a string.")
+        python_expr = translate(self.source)
+        object.__setattr__(self, "_python", python_expr)
+        object.__setattr__(self, "_code", compile_safely(python_expr))
+
+    def evaluate(
+        self,
+        ctx: Optional[Mapping[str, Any]] = None,
+        env: Optional[Mapping[str, Any]] = None,
+        extras: Optional[Mapping[str, Any]] = None,
+    ) -> Any:
+        scope = build_eval_scope(ctx or {}, env or {}, extras or {})
+        try:
+            return eval(self._code, {"__builtins__": {}}, scope)  # noqa: S307
+        except Exception as exc:
+            raise ExpressionEvaluationError(
+                f"Failed to evaluate expression '{self.source}': {exc}"
+            ) from exc
+
+
+def ensure_expression(value: Union[str, Expression]) -> Expression:
+    if isinstance(value, Expression):
+        return value
+    if isinstance(value, str):
+        return Expression(value.strip())
+    raise TypeError(f"Unsupported expression value: {value!r}")
+
+
+def translate(expr: str) -> str:
+    if "_ctx" in expr or "_env" in expr:
+        raise ExpressionSyntaxError("Use $ctx/$env references instead of _ctx/_env.")
+    def replace_token(match: re.Match[str]) -> str:
+        token = match.group(0)
+        if token.startswith("$ctx."):
+            path = token[len("$ctx.") :]
+            return f"_ctx('{path}')"
+        if token.startswith("$env."):
+            path = token[len("$env.") :]
+            return f"_env('{path}')"
+        raise ExpressionSyntaxError(f"Unsupported token '{token}'")
+
+    token_re = re.compile(r"\$(?:ctx|env)\.[A-Za-z_][\w.]*")
+    translated = token_re.sub(replace_token, expr.strip())
+    if "$" in translated:
+        raise ExpressionSyntaxError("All references must use $ctx.xxx or $env.xxx form.")
+    return translated
+
+
+ALLOWED_NODES = {
+    ast.Expression,
+    ast.BoolOp,
+    ast.BinOp,
+    ast.UnaryOp,
+    ast.Compare,
+    ast.And,
+    ast.Or,
+    ast.Not,
+    ast.Eq,
+    ast.NotEq,
+    ast.Gt,
+    ast.GtE,
+    ast.Lt,
+    ast.LtE,
+    ast.In,
+    ast.NotIn,
+    ast.Add,
+    ast.Sub,
+    ast.Mult,
+    ast.Div,
+    ast.Mod,
+    ast.Pow,
+    ast.USub,
+    ast.UAdd,
+    ast.Constant,
+    ast.Name,
+    ast.Load,
+    ast.Call,
+}
+
+
+def compile_safely(python_expr: str):
+    try:
+        tree = ast.parse(python_expr, mode="eval")
+    except SyntaxError as exc:
+        raise ExpressionSyntaxError(str(exc)) from exc
+
+    for node in ast.walk(tree):
+        if not isinstance(node, tuple(ALLOWED_NODES)):
+            raise ExpressionSyntaxError(f"Unsupported syntax: {type(node).__name__}")
+        if isinstance(node, ast.Name) and node.id not in {"_ctx", "_env"}:
+            raise ExpressionSyntaxError(f"Unknown identifier '{node.id}' in expression.")
+        if isinstance(node, ast.Call):
+            if not isinstance(node.func, ast.Name) or node.func.id not in {"_ctx", "_env"}:
+                raise ExpressionSyntaxError("Only $ctx/$env references are allowed.")
+            if len(node.args) != 1 or not isinstance(node.args[0], ast.Constant):
+                raise ExpressionSyntaxError("Context references must be constant strings.")
+    return compile(tree, "<expression>", "eval")
+
+
+def build_eval_scope(
+    ctx: Mapping[str, Any], env: Mapping[str, Any], extras: Mapping[str, Any]
+) -> Dict[str, Callable[[str], Any]]:
+    scope = {
+        "_ctx": lambda path: resolve_path(ctx, path, "$ctx"),
+        "_env": lambda path: resolve_path(env, path, "$env"),
+    }
+    scope.update(extras)
+    return scope
+
+
+def resolve_path(data: Mapping[str, Any], path: str, root: str):
+    if not DOT_PATH_RE.match(path):
+        raise ExpressionEvaluationError(f"Invalid path '{path}' for {root}.")
+    parts = path.split(".")
+    current: Any = data
+    for part in parts:
+        if isinstance(current, Mapping):
+            if part not in current:
+                raise ExpressionEvaluationError(f"{root}.{path} not found.")
+            current = current[part]
+        else:
+            if not hasattr(current, part):
+                raise ExpressionEvaluationError(f"{root}.{path} not found.")
+            current = getattr(current, part)
+    return current
+
+
+__all__ = [
+    "Expression",
+    "ensure_expression",
+    "ExpressionSyntaxError",
+    "ExpressionEvaluationError",
+]

pyoco/dsl/nodes.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Optional, Sequence as TypingSequence, Union
+
+from ..core.models import Task
+from .expressions import Expression, ensure_expression
+
+
+DEFAULT_CASE_VALUE = "__default__"
+
+
+class DSLNode:
+    """Base class for all DSL AST nodes."""
+
+
+@dataclass
+class TaskNode(DSLNode):
+    task: Task
+
+
+@dataclass
+class SubFlowNode(DSLNode):
+    steps: List[DSLNode] = field(default_factory=list)
+
+
+@dataclass
+class RepeatNode(DSLNode):
+    body: SubFlowNode
+    count: Union[int, Expression]
+
+
+@dataclass
+class ForEachNode(DSLNode):
+    body: SubFlowNode
+    source: Expression
+    alias: Optional[str] = None
+
+
+@dataclass
+class UntilNode(DSLNode):
+    body: SubFlowNode
+    condition: Expression
+    max_iter: Optional[int] = None
+
+
+@dataclass
+class CaseNode(DSLNode):
+    value: Union[str, int, float, bool]
+    target: SubFlowNode
+
+
+@dataclass
+class SwitchNode(DSLNode):
+    expression: Expression
+    cases: List[CaseNode] = field(default_factory=list)

pyoco/dsl/syntax.py

@@ -1,122 +1,268 @@
-from typing import Callable, Union, List, Tuple
-from ..core.models import Task, Flow
-from ..core import engine
+from __future__ import annotations
 
-# Global context
-_active_flow: Flow = None
+from dataclasses import dataclass
+from typing import Callable, Iterable, List, Sequence, Tuple, Union
+
+from ..core.models import Task
+from .expressions import Expression, ensure_expression
+from .nodes import (
+    CaseNode,
+    DSLNode,
+    ForEachNode,
+    RepeatNode,
+    SubFlowNode,
+    SwitchNode,
+    TaskNode,
+    UntilNode,
+    DEFAULT_CASE_VALUE,
+)
+RESERVED_CTX_KEYS = {"params", "results", "scratch", "loop", "loops", "env", "artifacts"}
+
+
+class FlowFragment:
+    """
+    Represents a fragment of a flow (sequence of DSL nodes). Every DSL
+    operator returns a FlowFragment so sub-flows can be composed before
+    being attached to a Flow.
+    """
+
+    __slots__ = ("_nodes",)
+
+    def __init__(self, nodes: Sequence[DSLNode]):
+        if not isinstance(nodes, (list, tuple)):
+            raise TypeError("FlowFragment expects a list/tuple of nodes.")
+        self._nodes: Tuple[DSLNode, ...] = tuple(nodes)
+
+    # Sequence composition -------------------------------------------------
+    def __rshift__(self, other: Union["FlowFragment", "TaskWrapper", Task]) -> "FlowFragment":
+        right = ensure_fragment(other)
+        self._link_to(right)
+        return FlowFragment(self._nodes + right._nodes)
+
+    # Loop support ---------------------------------------------------------
+    def __getitem__(self, selector: Union[int, str, Expression]) -> "FlowFragment":
+        """
+        Implements [] operator for repeat / for-each loops.
+        """
+
+        body = self.to_subflow()
+        if isinstance(selector, int):
+            if selector < 0:
+                raise ValueError("Repeat count must be non-negative.")
+            node = RepeatNode(body=body, count=selector)
+            return FlowFragment([node])
+
+        if isinstance(selector, Expression):
+            node = RepeatNode(body=body, count=selector)
+            return FlowFragment([node])
+
+        if isinstance(selector, str):
+            source, alias = parse_foreach_selector(selector)
+            node = ForEachNode(body=body, source=ensure_expression(source), alias=alias)
+            return FlowFragment([node])
+
+        raise TypeError(f"Unsupported loop selector: {selector!r}")
+
+    def __mod__(self, value: Union[str, Expression, Tuple[Union[str, Expression], int]]) -> "FlowFragment":
+        """
+        Implements the % operator for until loops.
+        """
+
+        max_iter = None
+        expr_value: Union[str, Expression]
+
+        if isinstance(value, tuple):
+            if len(value) != 2:
+                raise ValueError("Until tuple selector must be (expression, max_iter).")
+            expr_value, max_iter = value
+        else:
+            expr_value = value
+
+        node = UntilNode(
+            body=self.to_subflow(),
+            condition=ensure_expression(expr_value),
+            max_iter=max_iter,
+        )
+        return FlowFragment([node])
+
+    # Switch/case ----------------------------------------------------------
+    def __rrshift__(self, other: Union[str, int, float, bool]) -> CaseNode:
+        """
+        Enables `"X" >> fragment` syntax by reversing the operands.
+        """
+
+        if isinstance(other, str) and other == "*":
+            value = DEFAULT_CASE_VALUE
+        else:
+            value = other
+        return CaseNode(value=value, target=self.to_subflow())
+
+    # Helpers --------------------------------------------------------------
+    def to_subflow(self) -> SubFlowNode:
+        return SubFlowNode(list(self._nodes))
+
+    def task_nodes(self) -> List[Task]:
+        tasks: List[Task] = []
+        for node in self._nodes:
+            tasks.extend(_collect_tasks(node))
+        return tasks
+
+    def _first_task(self) -> Task | None:
+        for node in self._nodes:
+            tasks = _collect_tasks(node)
+            if tasks:
+                return tasks[0]
+        return None
+
+    def _last_task(self) -> Task | None:
+        for node in reversed(self._nodes):
+            tasks = _collect_tasks(node)
+            if tasks:
+                return tasks[-1]
+        return None
+
+    def _link_to(self, other: "FlowFragment"):
+        left_task = self._last_task()
+        right_task = other._first_task()
+        if left_task and right_task and left_task is not right_task:
+            right_task.dependencies.add(left_task)
+            left_task.dependents.add(right_task)
+
+    def has_control_flow(self) -> bool:
+        return any(not isinstance(node, TaskNode) for node in self._nodes)
+
+
+class TaskWrapper(FlowFragment):
+    """
+    Wraps a Task to handle DSL operators, while exposing the underlying
+    task for legacy access (e.g., `task.task.inputs = ...`).
+    """
+
+    __slots__ = ("task",)
+
+    def __init__(self, task: Task):
+        self.task = task
+        super().__init__([TaskNode(task)])
+
+    def __call__(self, *args, **kwargs) -> "TaskWrapper":
+        return self
+
+    def __and__(self, other):
+        return Parallel([self, other])
+
+    def __or__(self, other):
+        return Branch([self, other])
 
-def task(func: Callable) -> Task:
-    t = Task(func=func, name=func.__name__)
-    return t
 
 class Branch(list):
-    """Represents a branch of tasks (OR-split/join logic placeholder)."""
+    """Represents `A | B` OR-branches (legacy)."""
+
     def __rshift__(self, other):
-        # (A | B) >> C
-        # C depends on A and B.
-        # AND C.trigger_policy = "ANY"
-        
-        targets = []
-        if hasattr(other, 'task'):
-            targets = [other.task]
-        elif isinstance(other, (list, tuple)):
-             for item in other:
-                if hasattr(item, 'task'):
-                    targets.append(item.task)
-        
+        targets = _collect_target_tasks(other)
         for target in targets:
             target.trigger_policy = "ANY"
             for source in self:
-                if hasattr(source, 'task'):
+                if hasattr(source, "task"):
                     target.dependencies.add(source.task)
                     source.task.dependents.add(target)
-        
         return other
 
+
 class Parallel(list):
-    """Represents a parallel group of tasks (AND-split/join)."""
+    """Represents `A & B` parallel branches (legacy)."""
+
     def __rshift__(self, other):
-        # (A & B) >> C
-        # C depends on A AND B.
-        
-        targets = []
-        if hasattr(other, 'task'):
-            targets = [other.task]
-        elif isinstance(other, (list, tuple)):
-             for item in other:
-                if hasattr(item, 'task'):
-                    targets.append(item.task)
-        
+        targets = _collect_target_tasks(other)
         for target in targets:
             for source in self:
-                if hasattr(source, 'task'):
+                if hasattr(source, "task"):
                     target.dependencies.add(source.task)
                     source.task.dependents.add(target)
-        
         return other
 
-class TaskWrapper:
-    """
-    Wraps a Task to handle DSL operators and registration.
-    """
-    def __init__(self, task: Task):
-        self.task = task
-    
-    def __call__(self, *args, **kwargs):
-        # In this new spec, calling a task might not be strictly necessary for registration
-        # if we assume tasks are added to flow explicitly or via >>
-        # But let's keep the pattern: calling it returns a wrapper that can be chained
-        # We might need to store args/kwargs if we want to support them
-        return self
 
-    def __rshift__(self, other):
-        # self >> other
-        if isinstance(other, TaskWrapper):
-            other.task.dependencies.add(self.task)
-            self.task.dependents.add(other.task)
-            return other
-        elif isinstance(other, (list, tuple)):
-            # self >> (A & B) or self >> (A | B)
-            # If it's a Branch (from |), does it imply something different?
-            # Spec says: "Update Flow to handle Branch >> Task (set trigger_policy=ANY)"
-            # But here we are doing Task >> Branch.
-            # Task >> (A | B) means Task triggers both A and B?
-            # Usually >> means "follows".
-            # A >> (B | C) -> A triggers B and C?
-            # Or does it mean B and C depend on A? Yes.
-            # The difference between & and | is usually how they JOIN later, or how they are triggered?
-            # In Airflow, >> [A, B] means A and B depend on upstream.
-            # If we have (A | B) >> C, then C depends on A OR B.
-            # So if 'other' is a Branch, we just add dependencies as usual.
-            # The "OR" logic is relevant when 'other' connects to downstream.
-            
-            for item in other:
-                if isinstance(item, TaskWrapper):
-                    item.task.dependencies.add(self.task)
-                    self.task.dependents.add(item.task)
-            return other
-        return other
+def switch(expression: Union[str, Expression]) -> "SwitchBuilder":
+    return SwitchBuilder(expression=ensure_expression(expression))
 
-    def __and__(self, other):
-        # self & other (Parallel)
-        return Parallel([self, other])
 
-    def __or__(self, other):
-        # self | other (Branch)
-        # Return a Branch object containing both
-        return Branch([self, other])
+@dataclass
+class SwitchBuilder:
+    expression: Expression
+
+    def __getitem__(self, cases: Union[CaseNode, Sequence[CaseNode]]) -> FlowFragment:
+        if isinstance(cases, CaseNode):
+            case_list = [cases]
+        elif isinstance(cases, Sequence):
+            case_list = list(cases)
+        else:
+            raise TypeError("switch()[...] expects CaseNode(s)")
+
+        if not case_list:
+            raise ValueError("switch() requires at least one case.")
+        return FlowFragment([SwitchNode(expression=self.expression, cases=case_list)])
+
+
+# Helper utilities ---------------------------------------------------------
+def ensure_fragment(value: Union[FlowFragment, TaskWrapper, Task]) -> FlowFragment:
+    if isinstance(value, FlowFragment):
+        return value
+    if isinstance(value, TaskWrapper):
+        return value
+    if hasattr(value, "task"):
+        return FlowFragment([TaskNode(value.task)])
+    if isinstance(value, Task):
+        return TaskWrapper(value)
+    raise TypeError(f"Cannot treat {value!r} as a flow fragment.")
+
+
+def parse_foreach_selector(selector: str) -> Tuple[str, Union[str, None]]:
+    token = selector.strip()
+    alias = None
+    if " as " in token:
+        expr, alias = token.split(" as ", 1)
+        token = expr.strip()
+        alias = alias.strip()
+        if not alias or not alias.isidentifier() or alias in RESERVED_CTX_KEYS:
+            raise ValueError(f"Invalid foreach alias '{alias}'.")
+
+    return token, alias
+
+
+def _collect_tasks(obj) -> List[Task]:
+    if isinstance(obj, TaskNode):
+        return [obj.task]
+    if isinstance(obj, SubFlowNode):
+        tasks: List[Task] = []
+        for step in obj.steps:
+            tasks.extend(_collect_tasks(step))
+        return tasks
+    if isinstance(obj, RepeatNode):
+        return _collect_tasks(obj.body)
+    if isinstance(obj, ForEachNode):
+        return _collect_tasks(obj.body)
+    if isinstance(obj, UntilNode):
+        return _collect_tasks(obj.body)
+    if isinstance(obj, SwitchNode):
+        tasks: List[Task] = []
+        for case in obj.cases:
+            tasks.extend(_collect_tasks(case.target))
+        return tasks
+    return []
+
+
+def _collect_target_tasks(other) -> List[Task]:
+    targets = []
+    if hasattr(other, "task"):
+        targets = [other.task]
+    elif isinstance(other, (list, tuple)):
+        for item in other:
+            if hasattr(item, "task"):
+                targets.append(item.task)
+    return targets
 
-# We need to adapt the DSL to match the spec:
-# @task
-# def A(ctx, x:int)->int: ...
-# flow = Flow() >> A >> (B & C)
 
-# So A, B, C must be usable in the expression.
-# The @task decorator should return something that supports >>, &, |
+def task(func: Callable) -> TaskWrapper:
+    return TaskWrapper(Task(func=func, name=func.__name__))
 
-def task_decorator(func: Callable):
-    t = Task(func=func, name=func.__name__)
-    return TaskWrapper(t)
 
-# Re-export as task
-task = task_decorator
+__all__ = ["task", "FlowFragment", "switch", "TaskWrapper", "Branch", "Parallel"]

pyoco/dsl/validator.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List
+
+from ..core.models import Flow
+from .nodes import (
+    CaseNode,
+    ForEachNode,
+    RepeatNode,
+    SubFlowNode,
+    SwitchNode,
+    TaskNode,
+    UntilNode,
+    DSLNode,
+    DEFAULT_CASE_VALUE,
+)
+
+
+@dataclass
+class ValidationReport:
+    warnings: List[str] = field(default_factory=list)
+    errors: List[str] = field(default_factory=list)
+
+    @property
+    def status(self) -> str:
+        if self.errors:
+            return "error"
+        if self.warnings:
+            return "warning"
+        return "ok"
+
+    def to_dict(self) -> dict:
+        return {
+            "status": self.status,
+            "warnings": list(self.warnings),
+            "errors": list(self.errors),
+        }
+
+
+class FlowValidator:
+    """
+    Traverses a Flow's SubFlow definition and produces warnings/errors for
+    problematic control-flow constructs (unbounded loops, duplicate cases, etc.).
+    """
+
+    def __init__(self, flow: Flow):
+        self.flow = flow
+        self.report = ValidationReport()
+
+    def validate(self) -> ValidationReport:
+        program = self.flow.build_program()
+        self._visit_subflow(program, "flow")
+        return self.report
+
+    # Traversal helpers --------------------------------------------------
+    def _visit_subflow(self, subflow: SubFlowNode, path: str):
+        for idx, node in enumerate(subflow.steps):
+            self._visit_node(node, f"{path}.step[{idx}]")
+
+    def _visit_node(self, node: DSLNode, path: str):
+        if isinstance(node, TaskNode):
+            return
+        if isinstance(node, RepeatNode):
+            self._visit_subflow(node.body, f"{path}.repeat")
+        elif isinstance(node, ForEachNode):
+            self._visit_subflow(node.body, f"{path}.foreach")
+        elif isinstance(node, UntilNode):
+            self._validate_until(node, path)
+            self._visit_subflow(node.body, f"{path}.until")
+        elif isinstance(node, SwitchNode):
+            self._validate_switch(node, path)
+        elif isinstance(node, SubFlowNode):
+            self._visit_subflow(node, path)
+        else:
+            self.report.errors.append(f"{path}: Unknown node type {type(node).__name__}")
+
+    # Validators ---------------------------------------------------------
+    def _validate_until(self, node: UntilNode, path: str):
+        if node.max_iter is None:
+            self.report.warnings.append(f"{path}: Until loop missing max_iter (defaults to 1000).")
+
+    def _validate_switch(self, node: SwitchNode, path: str):
+        seen_values = set()
+        default_count = 0
+        for idx, case in enumerate(node.cases):
+            case_path = f"{path}.case[{idx}]"
+            if case.value == DEFAULT_CASE_VALUE:
+                default_count += 1
+                if default_count > 1:
+                    self.report.errors.append(f"{case_path}: Multiple default (*) cases are not allowed.")
+            else:
+                try:
+                    key = case.value
+                    if key in seen_values:
+                        self.report.errors.append(f"{case_path}: Duplicate switch value '{case.value}'.")
+                    else:
+                        seen_values.add(key)
+                except TypeError:
+                    self.report.errors.append(f"{case_path}: Unhashable switch value '{case.value}'.")
+            self._visit_subflow(case.target, f"{case_path}.target")
+
+        if default_count == 0:
+            self.report.warnings.append(f"{path}: Switch has no default (*) case.")