From b553d50932e27c9c6fdc9b462b3d3bef3f979415 Mon Sep 17 00:00:00 2001 From: Kristian Bendiksen Date: Mon, 26 Oct 2026 09:14:26 +0100 Subject: [PATCH 1/4] Add conditional branch execution (#16) Add SwitchTask/ConditionTask, when= guards and first_of() joins so a workflow can choose a branch at run time. Tasks whose guard does not match are skipped (TaskStatus.SKIPPED, on_task_skip hook); skips propagate, collect() drops skipped members, and a skipped result task fails the job with ResultTaskSkippedError. Supported in the builder API, YAML (when:, first_of:), Mermaid (switch diamonds, dashed guard edges), describe JSON and the editable definition round-trip. Adds conditional_if and conditional_switch examples and documentation. --- AGENTS.md | 6 +- README.md | 90 +++- docs/conditional-branches-plan.md | 83 ++++ examples/conditional_if/input.yaml | 7 + examples/conditional_if/pipeline.py | 198 ++++++++ examples/conditional_if/workflow.yaml | 29 ++ examples/conditional_switch/input.yaml | 8 + examples/conditional_switch/pipeline.py | 225 +++++++++ examples/conditional_switch/workflow.yaml | 37 ++ taskmaestro/__init__.py | 11 +- taskmaestro/conditions.py | 211 ++++++++ taskmaestro/definition.py | 128 ++++- taskmaestro/dependencies.py | 82 +++- taskmaestro/exceptions.py | 4 + taskmaestro/hooks/base.py | 5 + taskmaestro/hooks/logging.py | 3 + taskmaestro/introspection.py | 26 +- taskmaestro/job.py | 9 + taskmaestro/runner.py | 103 +++- taskmaestro/visualization.py | 73 ++- taskmaestro/workflow.py | 286 ++++++++++- taskmaestro/yaml_config.py | 112 ++++- tests/test_conditions.py | 557 ++++++++++++++++++++++ 23 files changed, 2248 insertions(+), 45 deletions(-) create mode 100644 docs/conditional-branches-plan.md create mode 100644 examples/conditional_if/input.yaml create mode 100644 examples/conditional_if/pipeline.py create mode 100644 examples/conditional_if/workflow.yaml create mode 100644 examples/conditional_switch/input.yaml create mode 100644 examples/conditional_switch/pipeline.py create mode 100644 examples/conditional_switch/workflow.yaml create mode 100644 taskmaestro/conditions.py create mode 100644 tests/test_conditions.py diff --git a/AGENTS.md b/AGENTS.md index 1b0b407..1f7b5a6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -23,7 +23,8 @@ mypy taskmaestro # type check (strict mode) | `taskmaestro/context.py` | `ExecutionContext` with correlation ID, logger, scratch dir, service registry | | `taskmaestro/task.py` | `Task[I, O]` ABC, type introspection (`get_input_type`, `get_output_type`) | | `taskmaestro/workflow.py` | `Workflow` (linear + DAG), `WorkflowBuilder`, validation (cycles, types, fan-in) | -| `taskmaestro/dependencies.py`, `taskmaestro/mapping.py` | Task handles, output references, `collect()`, mapped-task configuration | +| `taskmaestro/dependencies.py`, `taskmaestro/mapping.py` | Task handles (incl. guard sugar `on`/`otherwise`/`if_true`/`if_false`), output references, `collect()`, `first_of()`, mapped-task configuration | +| `taskmaestro/conditions.py` | `SwitchTask[I, L]`/`ConditionTask[I]` (label in a `Branch` output), `Condition` (unresolved `when=`), `Guard` (resolved), `when()` | | `taskmaestro/workflow_task.py` | Nested workflows wrapped as tasks | | `taskmaestro/yaml_config.py` | YAML parsing, task imports, workflow and input validation; `load_registered_workflow` pairs an entry-point workflow with an input file | | `taskmaestro/cli.py`, `taskmaestro/discovery.py` | CLI (`run`, `validate`, `graph`, `tasks list/describe`, `workflow list/describe`; workflow arguments accept a YAML path or a registered workflow name), plugin entry-point discovery (`load_registry` collects broken entry points instead of raising) | @@ -43,7 +44,8 @@ mypy taskmaestro # type check (strict mode) - **Timeouts**: `signal.setitimer`/`SIGALRM` (Unix only, main thread); gracefully warns if unavailable. All armed timeouts go through the process-wide `_AlarmScheduler`, which always arms the nearest expiry, so nested runs (`workflow_task`) cannot cancel or extend an enclosing run's timeouts - **Hook error swallowing**: `_emit()` wraps each hook call in try/except, reports via `warnings.warn(..., HookError, source=exc)` — message includes `repr(exc)`; `HookError` subclasses `UserWarning` so it can be filtered or escalated - **Inner-workflow failures**: `workflow_task` raises `WorkflowTaskError` (a `TaskExecutionError`) carrying `inner_job` and chaining the original exception via `__cause__`; `Job.exception` keeps the raw exception alongside `Job.error` -- **Validation order**: unique names → acyclic (DFS) → type chain → result task detection +- **Conditional branches**: a `when=` guard is a control edge (ordering + predicate, no data); guard sources count as upstreams for topo order, cycles and sinks. A task that is not run gets `TaskStatus.SKIPPED` + `skip_reason` and `on_task_skip(job, task_name, reason)` (no `task_start`). Skips propagate through plain deps; `collect()` drops skipped members; `first_of` takes the first member that ran; a skipped result task fails the job with `ResultTaskSkippedError`. Every switch label needs an `on()` or `otherwise()` +- **Validation order**: unique names → references → acyclic (DFS) → task maps → guards (incl. switch label coverage) → type chain (incl. `first_of`) → result task detection ## Testing Conventions diff --git a/README.md b/README.md index 99f2521..30fe49e 100644 --- a/README.md +++ b/README.md @@ -423,6 +423,89 @@ with `MappedOutput(root={})`. Per-item records are available in built-in logging, timing, and persistence hooks observe individual items. Concurrent mapped execution is intentionally deferred. +## Conditional Branches + +A workflow can choose between branches at run time. A *switch task* returns a +branch label; downstream tasks are guarded with `when=` and run only when their +guard matches. Tasks whose guard does not match are **skipped**, not failed. + +```python +from typing import Literal + +from taskmaestro import ConditionTask, SwitchTask, Workflow, first_of + + +class ChooseShipping(SwitchTask[Parcel, Literal["express", "freight", "standard"]]): + def decide(self, input: Parcel, ctx: ExecutionContext) -> Literal["express", "freight", "standard"]: + return "freight" if input.weight_kg > 100 else "standard" + + +builder = Workflow.builder("shipping") +parcel = builder.task(LoadParcel) +method = builder.task(ChooseShipping, depends_on=parcel) +express = builder.task(ExpressShipping, depends_on=parcel, when=method.on("express")) +freight = builder.task(FreightShipping, depends_on=parcel, when=method.on("freight")) +standard = builder.task(StandardShipping, depends_on=parcel, when=method.otherwise()) +builder.task(PrintLabel, parcel=parcel, shipment=first_of(express, freight, standard)) +``` + +- `SwitchTask[I, L]`: implement `decide()` and return a label of type `L`, a + `Literal[...]` or an `Enum`. The task output is a `Branch` model with a + `branch` field. +- `ConditionTask[I]` is a switch with the labels `True`/`False`; guard with + `handle.if_true()` and `handle.if_false()`. +- `handle.on(*labels)` runs a task for any of the labels. `handle.otherwise()` + runs it for every label that no `on()` guard claims. +- Any output field can be a guard too: `when(handle.field("kind"), equals=...)`, + `in_=[...]` or `truthy=True|False` (truthy/falsy needs a `bool` or optional + field; a non-bool value counts as truthy when it is not `None`, so `0` or `""` + still count as set). Guard values are validated strictly against the field + type: `"no"` or `1` is not accepted for a `bool` field, while the string value + of an enum member is. Shorthands: `handle.field("kind").equals(...)`, `.is_in(...)`, + `.truthy()`, `.falsy()`. +- `first_of(a, b, ...)` joins branches: the input gets the output of the first + listed member that ran. + +Validation happens at build time. Labels must belong to the switch's label +type, so typos are caught. Every label needs a guarded task or an +`otherwise()`. An `otherwise()` that can never match is an error. A switch with +no guarded tasks is an ordinary task. + +At run time: + +- A task is skipped when its guard does not match, when the guard's source was + skipped, when any plain dependency was skipped, or when every `first_of` + member was skipped. Skips propagate downstream. +- `collect()` leaves out skipped members. +- Skipped tasks get a `TaskResult` with `status=TaskStatus.SKIPPED` and a + `skip_reason`, and are listed in `job.skipped_tasks`. Hooks receive + `on_task_skip(job, task_name, reason)`; `on_task_start` is not called for + them. +- If the result task is skipped, the job fails with `ResultTaskSkippedError`. + +In YAML: + +```yaml +- task: pipeline.ExpressShipping + depends_on: pipeline.LoadParcel + when: {switch: pipeline.ChooseShipping, case: express} # case may be a list +- task: pipeline.StandardShipping + depends_on: pipeline.LoadParcel + when: {switch: pipeline.ChooseShipping, otherwise: true} +- task: pipeline.Audit + depends_on: pipeline.LoadParcel + when: {ref: [pipeline.LoadParcel, priority], truthy: true} # or equals: / in: [...] +- task: pipeline.PrintLabel + depends_on: + parcel: pipeline.LoadParcel + shipment: + first_of: [pipeline.ExpressShipping, pipeline.StandardShipping] +``` + +In Mermaid diagrams, switch tasks are diamonds and guards are dashed edges +labeled with their case (`express`, `otherwise`, `true`, `.kind = a`, ...). +`first_of` joins are drawn as rounded nodes. + ## ObjectModel `ObjectModel[T]` wraps arbitrary (non-Pydantic) objects so they can flow through workflows. Use it as a type alias for simple wrappers, or subclass it to add extra fields: @@ -616,7 +699,7 @@ graph TD build_report -->|AnalysisReport| _end_ ``` -Edges are labeled with data types. Fan-in edges show field names, and field routing edges show `.field: Type`. When a `JobConfiguration` is provided, configured tasks get dashed edges from a `JobConfiguration` node. +Edges are labeled with data types. Fan-in edges show field names, and field routing edges show `.field: Type`. When a `JobConfiguration` is provided, configured tasks get dashed edges from a `JobConfiguration` node. Switch tasks are drawn as diamonds with dashed, labeled guard edges (see [Conditional Branches](#conditional-branches)). ## Error Handling @@ -630,6 +713,7 @@ WorkflowRunnerError (base) │ └── WorkflowLookupError # Unknown or broken registered workflow name └── TaskExecutionError # Runtime task failure ├── MappedTaskExecutionError # One or more mapped items failed + ├── ResultTaskSkippedError # The result task was skipped by a condition ├── TaskOutputTypeError # Output type mismatch └── TaskTimeoutError # Task exceeded timeout ``` @@ -738,7 +822,7 @@ prints indented JSON. ## Examples -Four full example pipelines are included in the `examples/` directory: +Six full example pipelines are included in the `examples/` directory: | Example | Features | |---|---| @@ -746,6 +830,8 @@ Four full example pipelines are included in the `examples/` directory: | `examples/resinsight/` | `ObjectModel[T]` for gRPC objects, `JobConfiguration` with per-task config, named task instances, `config_fields`, YAML config | | `examples/image_processing/` | Nested workflows through `Workflow.as_task()` and YAML `workflow:`, typed boundaries, expanded Mermaid subgraph | | `examples/release_pipeline/` | Keyed `collect()` dependencies, mapped tasks, mapped output routing, per-task YAML config | +| `examples/conditional_if/` | If/else with `ConditionTask`, `if_true()`/`if_false()` guards, `first_of` join, YAML `when:` | +| `examples/conditional_switch/` | Multi-way `SwitchTask` with `on(...)` and `otherwise()`, `first_of` join, YAML `when:` | Run an example: diff --git a/docs/conditional-branches-plan.md b/docs/conditional-branches-plan.md new file mode 100644 index 0000000..5461bdf --- /dev/null +++ b/docs/conditional-branches-plan.md @@ -0,0 +1,83 @@ +# Conditional branch execution (issue #16) + +## Problem +The runner executes every task in topological order and stops at the first failure. Nothing lets a workflow pick a path at runtime. Issue #16 asks for conditional switching between branches. + +## Agreed design + +### Concepts +- **`when` guard**: a *control edge* on a task. It adds ordering and a runtime predicate, but it carries no data. The guarded task still gets its input through `depends_on`. +- **`SwitchTask[I, L]`**: the user implements `decide(input, ctx) -> L`, where `L` is a `Literal[...]` or `StrEnum`. The base class's final `run()` wraps the returned label in `Branch[L]`, a generic Pydantic model with one field, `branch: L`. `__init_subclass__` sets `Outputs = Branch[L]`, and the existing `get_output_type` honours that, so introspection does not change. The label set comes from `L`. +- **`ConditionTask[I]`**: a `SwitchTask` with branches `True`/`False`. `decide -> bool`. +- **Handle sugar**: `handle.on(*labels)` returns a guard (`in`). `handle.otherwise()` matches any label that no other guard on the same switch claims. For a ConditionTask: `on(True)` and `on(False)`, plus `if_true()`/`if_false()` aliases. +- **Generic field guards**: `when(ref, equals=…)`, `when(ref, in_=[…])`, `when(ref, truthy=True|False)` on any upstream output field. `ref` is an `OutputHandle`, a `(task, field)` tuple, or a name string. +- **Join**: `first_of(a, b, …)` is a dependency value like `collect()`. It resolves to the first member, in declaration order, that completed. +- Labels are chosen over target task ids. They are validated at build time, one label can enable many tasks, and renaming a task does not break routing. + +### Runtime semantics +- New `TaskStatus.SKIPPED`. A skipped task gets a `TaskResult(status=SKIPPED, output=None, error=None)` plus a reason, and `Event.TASK_SKIP` / `on_task_skip(job, task, reason)` is emitted. A skipped task does **not** emit `task_start`. +- A task is skipped when any of these holds: + 1. Its guard evaluates false. + 2. The guard's source task was skipped. + 3. Any plain data dependency (whole output, field ref, or a direct fan-in value) was skipped. + 4. Every member of one of its `first_of` dependencies was skipped. +- `collect()` drops skipped members. A list gets shorter; a dict loses those keys. If all members are skipped, the result is an empty list or dict, which is not a skip. +- Skips propagate transitively because tasks run in topological order. +- If the result task is skipped, the job **fails**. The new `ResultTaskSkippedError` goes in `job.exception`/`job.error`, `failed_task = result task`, and `JOB_FAIL` is emitted. +- If a guarded task is mapped and its guard is false, the whole node is skipped. No map items run. +- Nested workflows: if the inner result is skipped, the inner job fails, which raises `WorkflowTaskError` through the existing path. + +### Static validation (`Workflow._validate`) +New step after references: `_validate_guards`. The order becomes: unique names → references (now including guard and `first_of` refs) → acyclic (including control edges) → task maps → guards → types → result task. +- The guard source task and field must exist. A guard cannot read a field on a mapped task's output (`MappedOutput`). +- `on()`/`equals`/`in` values must be compatible with the field annotation. When the field is a `Literal` or enum, a value outside it is an error, which catches typos. +- `truthy`/`falsy` guards are only allowed on `bool` or Optional fields. +- For each switch: every declared label is consumed by some `on()` or covered by `otherwise()`, else build error. `otherwise()` matches the labels not claimed by any `on()` on that switch. Several tasks may use `otherwise()`, and they all mean the same thing. If every label is already claimed, `otherwise()` is an error because it is unreachable. +- `first_of` needs at least 2 members, and each member type must be compatible with the destination field (reuses `_is_type_compatible`). +- `topological_order`, `_find_sinks`, and cycle detection include guard sources as upstreams. A switch that is consumed only by guards is therefore not a sink. + +### Stored representation (`workflow.py`) +- `Workflow._guards: dict[str, Guard]`, where `Guard` is a frozen dataclass in a new `taskmaestro/conditions.py`: `ref: OutputRef`, `op: Literal["in","truthy","falsy","otherwise"]`, `values: tuple[Any, ...]`. `equals` is normalized to `in` with one value. The user-facing helpers `when()` and `first_of()` and the classes `SwitchTask`, `ConditionTask`, `Branch` also live there. +- `FirstOfRef(members: tuple[OutputRef, ...])` is stored as a `FanInValue`, next to `CollectionRef`. Update `_extract_upstream_names`, `_validate_types`, and `build()` copying. +- New public accessors: `get_guard(name)`, `is_conditional(name)`. + +## Surfaces to update +| Area | Change | +|---|---| +| `conditions.py` (new) | `Branch`, `SwitchTask`, `ConditionTask`, `Guard`, `Condition` (unresolved), `when()`; `first_of()`/`FirstOfRef` live in `dependencies.py` | +| `dependencies.py` | `TaskHandle.on()`, `.otherwise()`, `.if_true()`, `.if_false()`; `OutputHandle.equals()/.is_in()/.truthy()/.falsy()` sugar | +| `workflow.py` | `when=` kwarg on `add_task`/`task`/`map_task`; resolution of guard and `first_of` refs; validation; topo/sinks/cycles; snapshot copy in `build()` | +| `runner.py` | evaluate guard and skip rules before `TASK_START`; skipped set; `first_of` resolution; `collect()` dropping skipped members; result-task-skipped failure | +| `job.py` | `TaskStatus.SKIPPED`; `TaskResult.skip_reason: str \| None = None`; `Job.skipped_tasks` convenience property | +| `exceptions.py` | `ResultTaskSkippedError(TaskExecutionError)` (check hierarchy) | +| `hooks/base.py` | `Event.TASK_SKIP`, `on_task_skip` on the Protocol and `BaseHook` | +| `hooks/logging.py`, `timing.py`, `persistence.py` | log the skip; timing records nothing (or 0); persistence writes nothing for skipped tasks | +| `yaml_config.py` | `TaskConfig.when`: `{switch: , case: