pipeline/branch.py

"""Branch — parallel fork/join step."""

from __future__ import annotations

import asyncio
import dataclasses
from concurrent.futures import ThreadPoolExecutor
from enum import Enum
from types import MappingProxyType
from typing import Callable

from .context import StepContext
from .errors import BranchError


class MergeStrategy(Enum):
    """Built-in merge strategies for Branch outputs."""

    RAISE_ON_CONFLICT = "raise_on_conflict"
    LAST_WRITE_WINS = "last_write_wins"
    NAMESPACED = "namespaced"


# ---------------------------------------------------------------------------
# Built-in merge functions
# ---------------------------------------------------------------------------


def _merge_raise_on_conflict(ctxs: list[StepContext]) -> StepContext:
    """Raise if any two branches wrote different values for the same field.



    Metadata is always merged (union across all branches; last writer wins

    within metadata — there is no named-field semantic there).



    Uses ``type(ctxs[0])`` so subclass fields are included in the comparison.

    """
    if len(ctxs) == 1:
        return ctxs[0]

    conflicts: set[str] = set()
    for f in dataclasses.fields(type(ctxs[0])):
        if f.name == "metadata":
            continue
        first_val = getattr(ctxs[0], f.name)
        if any(getattr(ctx, f.name) != first_val for ctx in ctxs[1:]):
            conflicts.add(f.name)

    if conflicts:
        raise ValueError(
            f"Branch outputs conflict on fields {conflicts!r}. "
            "Use a different merge strategy or ensure branches write disjoint fields."
        )

    merged_meta: dict = {}
    for ctx in ctxs:
        merged_meta.update(ctx.metadata)

    return dataclasses.replace(ctxs[0], metadata=MappingProxyType(merged_meta))


def _merge_last_write_wins(ctxs: list[StepContext]) -> StepContext:
    """Last branch's value wins for every conflicting field.



    Uses ``type(ctxs[0])`` so subclass fields are included in the comparison.

    """
    if len(ctxs) == 1:
        return ctxs[0]

    # Start from first context, overlay with each subsequent one
    result = ctxs[0]
    ctx_type = type(ctxs[0])
    for ctx in ctxs[1:]:
        changes: dict = {}
        for f in dataclasses.fields(ctx_type):
            if f.name == "metadata":
                continue
            val = getattr(ctx, f.name)
            if val != getattr(result, f.name):
                changes[f.name] = val
        if changes:
            result = dataclasses.replace(result, **changes)

    merged_meta: dict = {}
    for ctx in ctxs:
        merged_meta.update(ctx.metadata)

    return dataclasses.replace(result, metadata=MappingProxyType(merged_meta))


def _merge_namespaced(ctxs: list[StepContext]) -> StepContext:
    """Each branch's output is stored at ``ctx.metadata["branch_N"]``.



    Named fields are taken from the first branch; no conflict is possible

    because branch outputs are kept in separate metadata keys.

    """
    base = ctxs[0]
    extra: dict = {f"branch_{i}": ctx for i, ctx in enumerate(ctxs)}
    merged_meta = MappingProxyType({**base.metadata, **extra})
    return dataclasses.replace(base, metadata=merged_meta)


_BUILTIN_MERGES: dict[MergeStrategy, Callable] = {
    MergeStrategy.RAISE_ON_CONFLICT: _merge_raise_on_conflict,
    MergeStrategy.LAST_WRITE_WINS: _merge_last_write_wins,
    MergeStrategy.NAMESPACED: _merge_namespaced,
}


# ---------------------------------------------------------------------------
# Branch
# ---------------------------------------------------------------------------


class Branch:
    """Runs multiple pipelines in parallel, then merges their outputs.



    ``Branch`` satisfies ``StepProtocol`` — it can be used wherever a step

    is expected.  ``requires`` and ``provides`` are inferred from the union

    of the child pipelines' contracts.



    In sync contexts (called directly), fan-out is via

    ``ThreadPoolExecutor``.  In async contexts (awaited), fan-out is via

    ``asyncio.gather``.



    All branches always run to completion before any failure is raised —

    ``BranchError`` carries the full list of failures.

    """

    def __init__(

        self,

        *pipelines: object,

        merge: MergeStrategy | Callable = MergeStrategy.RAISE_ON_CONFLICT,

    ) -> None:
        if not pipelines:
            raise ValueError("Branch requires at least one child pipeline.")

        self.pipelines = list(pipelines)

        if callable(merge) and not isinstance(merge, MergeStrategy):
            self._merge_fn: Callable = merge
        else:
            self._merge_fn = _BUILTIN_MERGES[merge]  # type: ignore[index]

        # Infer requires/provides from the union of child contracts
        all_requires: set[str] = set()
        all_provides: set[str] = set()
        for p in self.pipelines:
            all_requires |= set(getattr(p, "requires", frozenset()))
            all_provides |= set(getattr(p, "provides", frozenset()))

        self.requires: frozenset[str] = frozenset(all_requires)
        self.provides: frozenset[str] = frozenset(all_provides)

    # ------------------------------------------------------------------
    # Sync execution
    # ------------------------------------------------------------------

    def __call__(self, ctx: StepContext) -> StepContext:
        """Sync fan-out via ThreadPoolExecutor.



        All branches receive the same (frozen) context — no copy needed.

        All branches run to completion before any failure is raised.

        """
        with ThreadPoolExecutor(max_workers=len(self.pipelines)) as executor:
            futures: list = [executor.submit(p, ctx) for p in self.pipelines]  # type: ignore[arg-type]
            results: list[StepContext] = []
            failures: list[BaseException] = []
            for f in futures:
                try:
                    results.append(f.result())
                except BaseException as exc:  # noqa: BLE001
                    failures.append(exc)

        if failures:
            raise BranchError(failures)

        return self._merge_fn(results)

    # ------------------------------------------------------------------
    # Async execution
    # ------------------------------------------------------------------

    async def __call_async__(self, ctx: StepContext) -> StepContext:
        """Async fan-out via asyncio.gather.



        ``return_exceptions=True`` guarantees all branches run to completion

        even when one fails; the full failure list is surfaced via

        ``BranchError``.



        Sync child pipelines are wrapped with ``asyncio.to_thread`` so they

        run in a thread pool rather than blocking the event loop.

        """

        async def _run_child(child: object) -> StepContext:
            if asyncio.iscoroutinefunction(getattr(child, "__call__", None)):
                return await child(ctx)  # type: ignore[operator]
            if hasattr(child, "__call_async__"):
                return await child.__call_async__(ctx)  # type: ignore[union-attr]
            # Sync callable — run in thread pool so it doesn't block the loop
            return await asyncio.to_thread(child, ctx)  # type: ignore[arg-type]

        raw = await asyncio.gather(
            *[_run_child(p) for p in self.pipelines],
            return_exceptions=True,
        )
        failures = [r for r in raw if isinstance(r, BaseException)]
        if failures:
            raise BranchError(failures)
        return self._merge_fn([r for r in raw if not isinstance(r, BaseException)])