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

pyoco/core/engine.py

@@ -1,9 +1,32 @@
 import time
+import io
+import sys
+import traceback
 from typing import Dict, Any, List, Set, Optional
+import contextlib
 from .models import Flow, Task, RunContext, TaskState, RunStatus
-from .context import Context
+from .context import Context, LoopFrame
+from .exceptions import UntilMaxIterationsExceeded
 from ..trace.backend import TraceBackend
 from ..trace.console import ConsoleTraceBackend
+from ..dsl.nodes import TaskNode, RepeatNode, ForEachNode, UntilNode, SwitchNode, DEFAULT_CASE_VALUE
+from ..dsl.expressions import Expression
+
+class TeeStream:
+    def __init__(self, original):
+        self.original = original
+        self.buffer = io.StringIO()
+
+    def write(self, data):
+        self.original.write(data)
+        self.buffer.write(data)
+        return len(data)
+
+    def flush(self):
+        self.original.flush()
+
+    def getvalue(self):
+        return self.buffer.getvalue()
 
 class Engine:
     """
@@ -44,16 +67,31 @@ class Engine:
             run_context = RunContext()
         
         run_ctx = run_context
+        run_ctx.flow_name = flow.name
+        run_ctx.params = params or {}
         
         # Initialize all tasks as PENDING
         for task in flow.tasks:
             run_ctx.tasks[task.name] = TaskState.PENDING
+            run_ctx.ensure_task_record(task.name)
             
         ctx = Context(params=params or {}, run_context=run_ctx)
         self.trace.on_flow_start(flow.name, run_id=run_ctx.run_id)
         
         # Register active run
         self.active_runs[run_ctx.run_id] = run_ctx
+
+        if flow.has_control_flow():
+            try:
+                program = flow.build_program()
+                self._execute_subflow(program, ctx)
+                run_ctx.status = RunStatus.COMPLETED
+            except Exception:
+                run_ctx.status = RunStatus.FAILED
+                run_ctx.end_time = time.time()
+                raise
+            run_ctx.end_time = time.time()
+            return ctx
         
         try:
             executed: Set[Task] = set()
@@ -264,12 +302,130 @@ class Engine:
         run_ctx.end_time = time.time()
         return ctx
 
+    def _execute_subflow(self, subflow, ctx: Context):
+        for node in subflow.steps:
+            self._execute_node(node, ctx)
+
+    def _execute_node(self, node, ctx: Context):
+        if isinstance(node, TaskNode):
+            self._execute_task(node.task, ctx)
+        elif isinstance(node, RepeatNode):
+            self._execute_repeat(node, ctx)
+        elif isinstance(node, ForEachNode):
+            self._execute_foreach(node, ctx)
+        elif isinstance(node, UntilNode):
+            self._execute_until(node, ctx)
+        elif isinstance(node, SwitchNode):
+            self._execute_switch(node, ctx)
+        else:
+            raise TypeError(f"Unknown node type: {type(node)}")
+
+    def _execute_repeat(self, node: RepeatNode, ctx: Context):
+        count_value = self._resolve_repeat_count(node.count, ctx)
+        for index in range(count_value):
+            frame = LoopFrame(name="repeat", type="repeat", index=index, iteration=index + 1, count=count_value)
+            ctx.push_loop(frame)
+            try:
+                self._execute_subflow(node.body, ctx)
+            finally:
+                ctx.pop_loop()
+
+    def _execute_foreach(self, node: ForEachNode, ctx: Context):
+        sequence = self._eval_expression(node.source, ctx)
+        if not isinstance(sequence, (list, tuple)):
+            raise TypeError("ForEach source must evaluate to a list or tuple.")
+
+        total = len(sequence)
+        label = node.alias or node.source.source
+        for index, item in enumerate(sequence):
+            frame = LoopFrame(
+                name=f"foreach:{label}",
+                type="foreach",
+                index=index,
+                iteration=index + 1,
+                count=total,
+                item=item,
+            )
+            ctx.push_loop(frame)
+            if node.alias:
+                ctx.set_var(node.alias, item)
+            try:
+                self._execute_subflow(node.body, ctx)
+            finally:
+                if node.alias:
+                    ctx.clear_var(node.alias)
+                ctx.pop_loop()
+
+    def _execute_until(self, node: UntilNode, ctx: Context):
+        max_iter = node.max_iter or 1000
+        iteration = 0
+        last_condition = None
+        while True:
+            iteration += 1
+            frame = LoopFrame(
+                name="until",
+                type="until",
+                index=iteration - 1,
+                iteration=iteration,
+                condition=last_condition,
+                count=max_iter,
+            )
+            ctx.push_loop(frame)
+            try:
+                self._execute_subflow(node.body, ctx)
+                condition_result = bool(self._eval_expression(node.condition, ctx))
+            finally:
+                ctx.pop_loop()
+
+            last_condition = condition_result
+            if condition_result:
+                break
+            if iteration >= max_iter:
+                raise UntilMaxIterationsExceeded(node.condition.source, max_iter)
+
+    def _execute_switch(self, node: SwitchNode, ctx: Context):
+        value = self._eval_expression(node.expression, ctx)
+        default_case = None
+        for case in node.cases:
+            if case.value == DEFAULT_CASE_VALUE:
+                if default_case is None:
+                    default_case = case
+                continue
+            if case.value == value:
+                self._execute_subflow(case.target, ctx)
+                return
+        if default_case:
+            self._execute_subflow(default_case.target, ctx)
+    def _resolve_repeat_count(self, count_value, ctx: Context) -> int:
+        if isinstance(count_value, Expression):
+            resolved = self._eval_expression(count_value, ctx)
+        else:
+            resolved = count_value
+        if not isinstance(resolved, int):
+            raise TypeError("Repeat count must evaluate to an integer.")
+        if resolved < 0:
+            raise ValueError("Repeat count cannot be negative.")
+        return resolved
+
+    def _eval_expression(self, expression, ctx: Context):
+        if isinstance(expression, Expression):
+            return expression.evaluate(ctx=ctx.expression_data(), env=ctx.env_data())
+        return expression
+
     def _execute_task(self, task: Task, ctx: Context):
         # Update state to RUNNING
         from .models import TaskState
+        run_ctx = ctx.run_context
         if ctx.run_context:
             ctx.run_context.tasks[task.name] = TaskState.RUNNING
-            
+            record = ctx.run_context.ensure_task_record(task.name)
+            record.state = TaskState.RUNNING
+            record.started_at = time.time()
+            record.error = None
+            record.traceback = None
+        else:
+            record = None
+
         self.trace.on_node_start(task.name)
         start_time = time.time()
         # Retry loop
@@ -301,8 +457,17 @@ class Engine:
                     elif param_name in ctx.results:
                         kwargs[param_name] = ctx.results[param_name]
                 
-                result = task.func(**kwargs)
+                if record:
+                    record.inputs = {k: v for k, v in kwargs.items() if k != "ctx"}
+
+                stdout_capture = TeeStream(sys.stdout)
+                stderr_capture = TeeStream(sys.stderr)
+                with contextlib.redirect_stdout(stdout_capture), contextlib.redirect_stderr(stderr_capture):
+                    result = task.func(**kwargs)
                 ctx.set_result(task.name, result)
+                if run_ctx:
+                    run_ctx.append_log(task.name, "stdout", stdout_capture.getvalue())
+                    run_ctx.append_log(task.name, "stderr", stderr_capture.getvalue())
                 
                 # Handle outputs saving
                 for target_path in task.outputs:
@@ -333,10 +498,24 @@ class Engine:
                 # Update state to SUCCEEDED
                 if ctx.run_context:
                     ctx.run_context.tasks[task.name] = TaskState.SUCCEEDED
+                    if record:
+                        record.state = TaskState.SUCCEEDED
+                        record.ended_at = time.time()
+                        record.duration_ms = (record.ended_at - record.started_at) * 1000
+                        record.output = result
                 
                 return # Success
                 
             except Exception as e:
+                if run_ctx:
+                    run_ctx.append_log(task.name, "stdout", stdout_capture.getvalue() if 'stdout_capture' in locals() else "")
+                    run_ctx.append_log(task.name, "stderr", stderr_capture.getvalue() if 'stderr_capture' in locals() else "")
+                if record:
+                    record.state = TaskState.FAILED
+                    record.ended_at = time.time()
+                    record.duration_ms = (record.ended_at - record.started_at) * 1000
+                    record.error = str(e)
+                    record.traceback = traceback.format_exc()
                 if retries_left > 0:
                     retries_left -= 1
                     # Log retry?

pyoco/core/exceptions.py

@@ -0,0 +1,15 @@
+class ControlFlowError(Exception):
+    """Base error for control flow execution issues."""
+
+
+class UntilMaxIterationsExceeded(ControlFlowError):
+    def __init__(self, expression: str, max_iter: int):
+        super().__init__(f"Until condition '{expression}' exceeded max_iter={max_iter}")
+        self.expression = expression
+        self.max_iter = max_iter
+
+
+class SwitchNoMatch(ControlFlowError):
+    def __init__(self, expression: str):
+        super().__init__(f"Switch expression '{expression}' did not match any case.")
+        self.expression = expression

pyoco/core/models.py

@@ -1,8 +1,9 @@
-from typing import Any, Callable, Dict, List, Optional, Set, Union, ForwardRef
+from typing import Any, Callable, Dict, List, Optional, Set, Union
 from dataclasses import dataclass, field
 from enum import Enum
 import time
 import uuid
+import json
 
 @dataclass
 class Task:
@@ -56,16 +57,89 @@ class RunStatus(Enum):
     CANCELLING = "CANCELLING"
     CANCELLED = "CANCELLED"
 
+@dataclass
+class TaskRecord:
+    state: TaskState = TaskState.PENDING
+    started_at: Optional[float] = None
+    ended_at: Optional[float] = None
+    duration_ms: Optional[float] = None
+    error: Optional[str] = None
+    traceback: Optional[str] = None
+    inputs: Dict[str, Any] = field(default_factory=dict)
+    output: Any = None
+    artifacts: Dict[str, Any] = field(default_factory=dict)
+
+
 @dataclass
 class RunContext:
     """
     Holds the state of a single workflow execution.
     """
     run_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    flow_name: str = "main"
+    params: Dict[str, Any] = field(default_factory=dict)
     status: RunStatus = RunStatus.RUNNING
     tasks: Dict[str, TaskState] = field(default_factory=dict)
+    task_records: Dict[str, TaskRecord] = field(default_factory=dict)
     start_time: float = field(default_factory=time.time)
     end_time: Optional[float] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    logs: List[Dict[str, Any]] = field(default_factory=list)
+    _pending_logs: List[Dict[str, Any]] = field(default_factory=list, repr=False)
+    _log_seq: int = field(default=0, repr=False)
+    log_bytes: Dict[str, int] = field(default_factory=dict)
+    metrics_recorded_tasks: Set[str] = field(default_factory=set, repr=False)
+    metrics_run_observed: bool = field(default=False, repr=False)
+    webhook_notified_status: Optional[str] = field(default=None, repr=False)
+
+    def ensure_task_record(self, task_name: str) -> TaskRecord:
+        if task_name not in self.task_records:
+            self.task_records[task_name] = TaskRecord()
+        return self.task_records[task_name]
+
+    def append_log(self, task_name: str, stream: str, payload: str):
+        if not payload:
+            return
+        entry = {
+            "seq": self._log_seq,
+            "task": task_name,
+            "stream": stream,
+            "text": payload,
+            "timestamp": time.time(),
+        }
+        self._log_seq += 1
+        self.logs.append(entry)
+        self._pending_logs.append(entry)
+
+    def drain_logs(self) -> List[Dict[str, Any]]:
+        drained = list(self._pending_logs)
+        self._pending_logs.clear()
+        return drained
+
+    def serialize_task_records(self) -> Dict[str, Any]:
+        serialized: Dict[str, Any] = {}
+        for name, record in self.task_records.items():
+            serialized[name] = {
+                "state": record.state.value if hasattr(record.state, "value") else record.state,
+                "started_at": record.started_at,
+                "ended_at": record.ended_at,
+                "duration_ms": record.duration_ms,
+                "error": record.error,
+                "traceback": record.traceback,
+                "inputs": {k: self._safe_value(v) for k, v in record.inputs.items()},
+                "output": self._safe_value(record.output),
+                "artifacts": record.artifacts,
+            }
+        return serialized
+
+    def _safe_value(self, value: Any) -> Any:
+        if isinstance(value, (str, int, float, bool)) or value is None:
+            return value
+        try:
+            json.dumps(value)
+            return value
+        except Exception:
+            return repr(value)
 
 @dataclass
 class Flow:
@@ -78,8 +152,27 @@ class Flow:
     name: str = "main"
     tasks: Set[Task] = field(default_factory=set)
     _tail: Set[Task] = field(default_factory=set)
+    _definition: List[Any] = field(default_factory=list, repr=False)
+    _has_control_flow: bool = False
     
     def __rshift__(self, other):
+        from ..dsl.syntax import TaskWrapper, FlowFragment, ensure_fragment
+
+        if isinstance(other, TaskWrapper):
+            fragment = other
+            self._record_fragment(fragment)
+            self._append_task(fragment.task)
+            return self
+
+        if hasattr(other, "to_subflow"):
+            fragment = other if isinstance(other, FlowFragment) else ensure_fragment(other)
+            self._record_fragment(fragment)
+            if not self._has_control_flow and not fragment.has_control_flow():
+                self._append_linear_fragment(fragment)
+            else:
+                self._has_control_flow = True
+            return self
+
         # Flow >> Task/List/Branch
         new_tasks = []
         is_branch = False
@@ -155,3 +248,39 @@ class Flow:
     def add_task(self, task: Task):
         self.tasks.add(task)
 
+    def has_control_flow(self) -> bool:
+        return self._has_control_flow
+
+    def build_program(self):
+        from ..dsl.nodes import SubFlowNode
+        return SubFlowNode(list(self._definition))
+
+    def _record_fragment(self, fragment):
+        from ..dsl.nodes import TaskNode
+        subflow = fragment.to_subflow()
+        self._definition.extend(subflow.steps)
+        for task in fragment.task_nodes():
+            self.add_task(task)
+        if any(not isinstance(step, TaskNode) for step in subflow.steps):
+            self._has_control_flow = True
+
+    def _append_linear_fragment(self, fragment):
+        subflow = fragment.to_subflow()
+        for step in subflow.steps:
+            if hasattr(step, "task"):
+                self._append_task(step.task)
+
+    def _append_task(self, task: Task):
+        self.add_task(task)
+        if self._has_control_flow:
+            self._tail = {task}
+            return
+
+        if not self._tail:
+            self._tail = {task}
+            return
+
+        for tail_task in self._tail:
+            tail_task.dependents.add(task)
+            task.dependencies.add(tail_task)
+        self._tail = {task}

pyoco/discovery/loader.py

@@ -1,9 +1,10 @@
 import importlib
 import pkgutil
 import sys
-from typing import Dict, List, Any
+from typing import Dict, List, Any, Set
 from ..core.models import Task
 from ..dsl.syntax import TaskWrapper
+from .plugins import PluginRegistry, iter_entry_points
 
 class TaskLoader:
     def __init__(self, config: Any, strict: bool = False):
@@ -11,6 +12,7 @@ class TaskLoader:
         self.strict = strict
         self.tasks: Dict[str, Task] = {}
         self._explicit_tasks: Set[str] = set()
+        self.plugin_reports: List[Dict[str, Any]] = []
 
     def load(self):
         # Load explicitly defined tasks in config FIRST (Higher priority)
@@ -31,6 +33,8 @@ class TaskLoader:
         for pattern in self.config.discovery.glob_modules:
             self._load_glob_modules(pattern)
 
+        self._load_entry_point_plugins()
+
     def _register_task(self, name: str, task: Task):
         if name in self.tasks:
             if name in self._explicit_tasks:
@@ -97,6 +101,33 @@ class TaskLoader:
             module_name = rel_path.replace(os.sep, ".")[:-3] # strip .py
             self._load_module(module_name)
 
+    def _load_entry_point_plugins(self):
+        entries = iter_entry_points()
+        for ep in entries:
+            info = {
+                "name": ep.name,
+                "value": ep.value,
+                "module": getattr(ep, "module", ""),
+                "tasks": [],
+                "warnings": [],
+            }
+            registry = PluginRegistry(self, ep.name)
+            try:
+                hook = ep.load()
+                if not callable(hook):
+                    raise TypeError("Entry point must be callable")
+                hook(registry)
+                info["tasks"] = list(registry.records)
+                info["warnings"] = list(registry.warnings)
+                if not registry.records:
+                    info["warnings"].append("no tasks registered")
+            except Exception as exc:
+                info["error"] = str(exc)
+                if self.strict:
+                    raise
+                print(f"Warning: Plugin '{ep.name}' failed to load: {exc}")
+            self.plugin_reports.append(info)
+
     def _scan_module(self, module: Any):
         for name, obj in vars(module).items():
             if isinstance(obj, TaskWrapper):

pyoco/discovery/plugins.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+from importlib import metadata as importlib_metadata
+from typing import Any, Callable, Dict, List, Optional, Type
+
+from ..core.models import Task
+from ..dsl.syntax import TaskWrapper
+
+
+class CallablePluginTask(Task):
+    """Lightweight subclass so callable registrations still appear as Task-derived."""
+
+    def __init__(self, func: Callable, name: str):
+        super().__init__(func=func, name=name)
+
+
+def iter_entry_points(group: str = "pyoco.tasks"):
+    eps = importlib_metadata.entry_points()
+    if hasattr(eps, "select"):
+        return list(eps.select(group=group))
+    return list(eps.get(group, []))
+
+
+def list_available_plugins() -> List[Dict[str, Any]]:
+    plugins = []
+    for ep in iter_entry_points():
+        plugins.append(
+            {
+                "name": ep.name,
+                "module": getattr(ep, "module", ""),
+                "value": ep.value,
+            }
+        )
+    return plugins
+
+
+class PluginRegistry:
+    def __init__(self, loader: Any, provider_name: str) -> None:
+        self.loader = loader
+        self.provider_name = provider_name
+        self.registered_names: List[str] = []
+        self.records: List[Dict[str, Any]] = []
+        self.warnings: List[str] = []
+
+    def task(
+        self,
+        func: Optional[Callable] = None,
+        *,
+        name: Optional[str] = None,
+        inputs: Optional[Dict[str, Any]] = None,
+        outputs: Optional[List[str]] = None,
+    ):
+        if func is not None:
+            self.register_callable(
+                func,
+                name=name,
+                inputs=inputs or {},
+                outputs=outputs or [],
+            )
+            return func
+
+        def decorator(inner: Callable):
+            self.register_callable(
+                inner,
+                name=name,
+                inputs=inputs or {},
+                outputs=outputs or [],
+            )
+            return inner
+
+        return decorator
+
+    def register_callable(
+        self,
+        func: Callable,
+        *,
+        name: Optional[str] = None,
+        inputs: Optional[Dict[str, Any]] = None,
+        outputs: Optional[List[str]] = None,
+    ) -> Task:
+        task_name = name or getattr(func, "__name__", f"{self.provider_name}_task")
+        task = CallablePluginTask(func=func, name=task_name)
+        if inputs:
+            task.inputs.update(inputs)
+        if outputs:
+            task.outputs.extend(outputs)
+        self._finalize_task(task, origin="callable")
+        return task
+
+    def task_class(
+        self,
+        task_cls: Type[Task],
+        *args: Any,
+        name: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Task:
+        if not issubclass(task_cls, Task):
+            raise TypeError(f"{task_cls} is not a Task subclass")
+        task = task_cls(*args, **kwargs)
+        if name:
+            task.name = name
+        self._finalize_task(task, origin="task_class")
+        return task
+
+    def add(self, obj: Any, *, name: Optional[str] = None) -> None:
+        if isinstance(obj, TaskWrapper):
+            task = obj.task
+            if name:
+                task.name = name
+            self._finalize_task(task, origin="wrapper")
+        elif isinstance(obj, Task):
+            if name:
+                obj.name = name
+            origin = "task_class" if obj.__class__ is not Task else "task"
+            self._finalize_task(obj, origin=origin)
+        elif callable(obj):
+            self.register_callable(obj, name=name)
+        else:
+            raise TypeError(f"Unsupported task object: {obj!r}")
+
+    def _finalize_task(self, task: Task, origin: str) -> None:
+        warnings = self._validate_task(task, origin)
+        self.loader._register_task(task.name, task)
+        self.registered_names.append(task.name)
+        self.records.append(
+            {
+                "name": task.name,
+                "origin": origin,
+                "class": task.__class__.__name__,
+                "warnings": warnings,
+            }
+        )
+        for msg in warnings:
+            self.warnings.append(f"{task.name}: {msg}")
+
+    def _validate_task(self, task: Task, origin: str) -> List[str]:
+        warnings: List[str] = []
+        if not getattr(task, "name", None):
+            generated = f"{self.provider_name}_{len(self.registered_names) + 1}"
+            task.name = generated
+            warnings.append(f"name missing; auto-assigned '{generated}'")
+        if not callable(getattr(task, "func", None)):
+            warnings.append("task.func is not callable")
+        if origin == "callable":
+            warnings.append("registered via callable; prefer Task subclass for extensibility")
+        if task.__class__ is Task and origin not in ("callable", "wrapper"):
+            warnings.append("plain Task instance detected; subclass Task for metadata support")
+        return warnings