Pipeline

def add_step[**P, OutputT](
    self, transform: Transform[P, OutputT], name: str | None = None
) -> Step[OutputT]:
    """Add a transform as a named step.

    Args:
        transform: Function decorated with ``@transform``.
        name: Step name. Defaults to the transform's function name.

    Returns:
        The new step.

    Raises:
        ValueError: If the name is already used or cannot be inferred.
        TypeError: If the transform has unsupported parameters or context
            dependencies that the pipeline cannot supply.
    """
    if not isinstance(transform, Transform):
        raise TypeError(
            "add_step() requires a function decorated with @transform; "
            f"got {type(transform).__name__}."
        )
    if name is None:
        name = transform.default_name
    if name is None:
        raise ValueError(
            "Cannot infer a name for this transform. Pass name='step_name' "
            "to add_step()."
        )
    if name in self._steps_by_name:
        raise ValueError(
            f"Step name '{name}' is already registered. Choose a unique name."
        )

    self._validate_transform_parameters(transform)
    self._validate_transform_context(transform)

    step = Step(
        name=name,
        transform=transform,
        _pipeline=self,
    )
    node_index = self._dag.add_node(step)
    self._node_indices[step] = node_index
    self._steps_by_name[name] = step
    return step

def connect(
    self, upstream: Step, downstream: Step, *, param: str | None = None
) -> None:
    """Connect an output to a compatible input on another step.

    Args:
        upstream: Step that produces the value.
        downstream: Step that receives the value.
        param: Downstream parameter name. Required when more than one
            input can accept the output.

    Raises:
        PipelineConnectionError: If the steps or their contracts cannot
            be connected.
        PipelineCycleError: If the connection would create a cycle.
    """
    for role, step in (("upstream", upstream), ("downstream", downstream)):
        if not isinstance(step, Step):
            raise TypeError(
                f"{role.capitalize()} must be a Step; got {type(step).__name__}."
            )
        if step._pipeline is not self:
            raise PipelineConnectionError(
                f"{role.capitalize()} step '{step.name}' belongs to another "
                "pipeline. Add it to this pipeline before connecting it."
            )

    downstream_spec = downstream.transform.__transform_spec__
    upstream_spec = upstream.transform.__transform_spec__

    upstream_contract = upstream_spec.output_contract
    if isinstance(upstream_contract, VoidContract):
        raise PipelineConnectionError(
            f"Upstream step '{upstream.name}' produces None. Declare an output "
            "contract on its transform or connect a step that produces a value."
        )

    compatible_specs: dict[str, Contract] = {}
    for param_name, input_contract in downstream_spec.input_contracts.items():
        if input_contract.accepts(upstream_contract):
            compatible_specs[param_name] = input_contract

    if len(compatible_specs) > 1:
        if param is None:
            names = ", ".join(sorted(compatible_specs))
            raise PipelineConnectionError(
                f"Step '{upstream.name}' produces {describe_contract(upstream_contract)}, "
                f"which matches multiple inputs on '{downstream.name}': {names}. "
                "Pass param='input_name' to connect()."
            )
        if param not in compatible_specs:
            names = ", ".join(sorted(compatible_specs))
            raise PipelineConnectionError(
                f"Input '{param}' on step '{downstream.name}' is not compatible "
                f"with '{upstream.name}' ({describe_contract(upstream_contract)}). "
                f"Choose one of: {names}."
            )
        downstream_parameter_name = param
    elif len(compatible_specs) == 1:
        compatible_parameter = next(iter(compatible_specs))
        if param is not None and param != compatible_parameter:
            raise PipelineConnectionError(
                f"Input '{param}' on step '{downstream.name}' does not match "
                f"the output of '{upstream.name}' ({describe_contract(upstream_contract)}). "
                f"Connect to the compatible input '{compatible_parameter}'."
            )
        downstream_parameter_name = compatible_parameter
    else:
        inputs = [
            f"    '{name}': {describe_contract(contract)}"
            for name, contract in downstream_spec.input_contracts.items()
        ] or ["    none"]
        inputs_text = "\n".join(inputs)
        requested = f"\n  Requested input: '{param}'" if param else ""
        action = (
            "Add an annotated input parameter to the downstream transform."
            if not downstream_spec.input_contracts
            else (
                "Change an input or output annotation so their DataFrame "
                "backends, types, and schema models match."
                if isinstance(upstream_contract, DataFrameContract)
                else "Change an input or output annotation so the produced "
                "value type matches an input."
            )
        )
        raise PipelineConnectionError(
            f"No compatible contract between steps '{upstream.name}' and "
            f"'{downstream.name}':\n"
            f"  Output of '{upstream.name}': {describe_contract(upstream_contract)}\n"
            f"  Inputs of '{downstream.name}':\n"
            f"{inputs_text}{requested}\n"
            f"{action}"
        )

    downstream_index = self._node_indices[downstream]
    for connected_upstream_index, _, connected_parameter_name in self._dag.in_edges(
        downstream_index
    ):
        if connected_parameter_name != downstream_parameter_name:
            continue
        if connected_upstream_index == self._node_indices[upstream]:
            return
        connected_upstream = self._dag[connected_upstream_index]
        raise PipelineConnectionError(
            f"Step '{connected_upstream.name}' is already connected to "
            f"'{downstream.name}' input '{downstream_parameter_name}'. "
            "Disconnect that input before connecting another step."
        )

    try:
        self._dag.add_edge(
            self._node_indices[upstream],
            downstream_index,
            downstream_parameter_name,
        )
    except rx.DAGWouldCycle:
        raise PipelineCycleError(
            f"Connecting step '{upstream.name}' to '{downstream.name}' would "
            "create a cycle. Remove a connection so the steps form a DAG."
        )

def connect_many(
    self,
    upstream_steps: Sequence[Step],
    downstream: Step,
) -> None:
    """Connect several upstream steps to one downstream step.

    Each output must match one unambiguous downstream input.

    Args:
        upstream_steps: Nonempty sequence of steps that provide inputs.
        downstream: Step that receives the inputs.

    Raises:
        ValueError: If ``upstream_steps`` is empty.
        PipelineConnectionError: If an output cannot be connected.
        PipelineCycleError: If a connection would create a cycle.
    """
    if not upstream_steps:
        raise ValueError(
            "connect_many requires at least one upstream step. Pass a nonempty "
            "sequence of Step instances."
        )
    for upstream in upstream_steps:
        self.connect(upstream, downstream)

def get_upstream_steps(self, step: Step) -> set[Step]:
    """Return the steps directly connected to ``step`` as inputs."""
    return set(self._dag.predecessors(self._node_indices[step]))

def get_steps_in_execution_order(self) -> list[Step]:
    """Return all steps in an order that respects their connections."""
    node_indices = rx.topological_sort(self._dag)
    return [self._dag[node_index] for node_index in node_indices]

def run(
    self,
    context: ContextT | None = None,
    *,
    transform_overrides: Mapping[str, Transform] | None = None,
    include_side_effects: bool = True,
) -> PipelineResult:
    """Invoke each step and return its recorded inputs and outputs.

    Args:
        context: Instance of the context dataclass, when the pipeline
            was created with one.
        transform_overrides: Replacement transforms keyed by step name.
            Each replacement must have compatible input contracts and the
            same output contract.
        include_side_effects: Whether to run transforms marked with
            ``side_effect=True``. Defaults to ``True``.

    Returns:
        A result containing an invocation record for every executed step.

    Raises:
        TypeError: If the context is missing or has the wrong type.
        PipelineOverrideError: If a replacement is invalid.
        PipelineExecutionError: If a required input is unconnected or a
            step fails.
    """
    if self._context_type is None:
        if context is not None:
            raise TypeError(
                "A context-free pipeline does not accept a context; "
                f"got {type(context).__name__}. Call run() without one."
            )
    elif context is None:
        raise TypeError(
            f"Pipeline[{self._context_type.__name__}] requires a context. "
            f"Pass a {self._context_type.__name__} instance to run()."
        )
    elif not isinstance(context, self._context_type):
        raise TypeError(
            f"expected context of type {self._context_type.__name__}, "
            f"got {type(context).__name__}. Pass a "
            f"{self._context_type.__name__} instance to run()."
        )

    overrides = self._resolve_transform_overrides(transform_overrides)
    steps_to_run = [
        step
        for step in self.get_steps_in_execution_order()
        if include_side_effects or not step.transform.side_effect
    ]
    self._validate_source_steps(overrides, steps_to_run)
    run_started_at = perf_counter()
    logger.info(
        "Pipeline.run() started (steps: %d, context: %s, overrides: %d)",
        len(self._node_indices),
        type(context).__name__ if context is not None else "none",
        len(overrides),
    )
    for name, replacement in overrides.items():
        logger.debug(
            "Override for step '%s': transform '%s'",
            name,
            replacement.default_name,
        )

    outputs: dict[str, Any] = {}
    step_invocations: list[StepInvocation[Any]] = []

    for step in steps_to_run:
        effective_transform = overrides.get(step.name, step.transform)
        spec = effective_transform.__transform_spec__
        arguments: dict[str, Any] = {}
        sources: dict[str, Step[Any] | ContextSource] = {}
        step_index = self._node_indices[step]

        for upstream_index, _, parameter_name in self._dag.in_edges(step_index):
            upstream_step = self._dag[upstream_index]
            arguments[parameter_name] = outputs[upstream_step.name]
            sources[parameter_name] = upstream_step
        for parameter_name, (
            param_type,
            marker,
        ) in spec.context_parameters.items():
            dependency_type = marker.type or param_type
            context_field = self._context_fields[dependency_type]
            arguments[parameter_name] = getattr(context, context_field)
            sources[parameter_name] = ContextSource(context_field)

        signature = effective_transform.get_signature()
        missing = {
            name
            for name in spec.input_contracts
            if name not in arguments
            and signature.parameters[name].default is inspect.Parameter.empty
        }
        if missing:
            names = ", ".join(sorted(missing))
            raise PipelineExecutionError(
                f"Pipeline step '{step.name}' has unconnected inputs: {names}. "
                "Connect an upstream step to each input."
            )

        logger.debug(
            "Invoking step '%s' with transform '%s' "
            "(inputs: %s; context fields: %s)",
            step.name,
            effective_transform.default_name,
            tuple(sorted(arguments)),
            tuple(
                sorted(
                    source.field
                    for source in sources.values()
                    if isinstance(source, ContextSource)
                )
            ),
        )
        step_started_at = perf_counter()
        try:
            bound = signature.bind(**arguments)
            bound.apply_defaults()
            result = effective_transform(**arguments)
        except Exception as error:
            if isinstance(error, PipelineExecutionError):
                raise
            raise PipelineExecutionError(
                f"Pipeline step '{step.name}' failed: "
                f"{type(error).__name__}: {error}"
            ) from error
        logger.debug(
            "Step '%s' returned in %.3fs (return type: %s)",
            step.name,
            perf_counter() - step_started_at,
            type(result).__name__,
        )

        if result is not None:
            outputs[step.name] = result

        step_invocations.append(
            StepInvocation(
                step=step,
                transform=effective_transform,
                inputs=bound_inputs(bound.arguments, sources),
                output=result,
            )
        )

    pipeline_result = PipelineResult(tuple(step_invocations))
    logger.info(
        "Pipeline.run() returned in %.3fs (step invocations: %d, outputs: %d)",
        perf_counter() - run_started_at,
        len(step_invocations),
        len(pipeline_result),
    )
    logger.debug(
        "Pipeline result step invocations: %s",
        ", ".join(
            invocation.step.name for invocation in pipeline_result.step_invocations
        )
        or "<none>",
    )
    return pipeline_result

def visualize(self, *, show: bool = True):
    """Draw the pipeline and return its Matplotlib figure.

    Args:
        show: Whether to display the figure. Set to ``False`` to save or
            customize it without opening a window.

    Raises:
        ImportError: If the ``vis`` extra is not installed.
    """
    if not _MATPLOTLIB_AVAILABLE:
        raise ImportError(
            "Pipeline visualization requires matplotlib; "
            "install it with `pip install 'data-joinery[vis]'`."
        )

    assert plt is not None
    figure = draw_pipeline(self._dag)
    if show:
        plt.show()
    return figure

def get_step_invocation(self, step: Step[Any] | str) -> StepInvocation[Any]:
    """Get a step's invocation record, including steps that return `None`.

    Args:
        step: A step from this pipeline or its name. A step object must be
            the same object that was invoked, not just have the same name.

    Returns:
        The step invocation's recorded inputs and output.

    Raises:
        KeyError: If the step name is absent or the step was not invoked.
    """
    if isinstance(step, str):
        return self._by_step_name[step]
    invocation = self._by_step_name.get(step.name)
    if invocation is None or invocation.step is not step:
        raise KeyError(f"Step '{step.name}' is not in this pipeline result.")
    return invocation

def get_output(
    self, step: Step[Any] | str, value_type: type[Any] | None = None
) -> Any:
    """Get a step's output.

    A `Step` keeps its declared output type. When using a step name, pass
    the expected runtime type.

    Args:
        step: A step from this pipeline or its name.
        value_type: Required runtime type when `step` is a name.

    Returns:
        The step's output, including `None` for a step that returns it.

    Raises:
        KeyError: If the step is absent.
        TypeError: If the type is missing for a name, supplied for a `Step`,
            or does not match the named output.
    """
    if not isinstance(step, str):
        if value_type is not None:
            raise TypeError("Pass value_type only when looking up a step by name.")
        return self.get_step_invocation(step).output
    if value_type is None:
        raise TypeError("Looking up an output by step name requires value_type.")
    value = self._by_step_name[step].output
    if not isinstance(value, value_type):
        raise TypeError(
            f"Output of step '{step}' is {type(value).__name__}, "
            f"expected {value_type.__name__}."
        )
    return value

def get_input[T](
    self, step: Step[Any] | str, parameter_name: str, value_type: type[T]
) -> T:
    """Get a value bound to a step's input parameter before input coercion.

    Args:
        step: A step from this pipeline or its name.
        parameter_name: Name of the transform parameter.
        value_type: Required runtime type of the bound value.

    Returns:
        The recorded input value.

    Raises:
        KeyError: If the step or parameter is absent.
        TypeError: If the value is not an instance of `value_type`.
    """
    invocation = self.get_step_invocation(step)
    binding = invocation.inputs[parameter_name]
    if not isinstance(binding.value, value_type):
        raise TypeError(
            f"Input '{parameter_name}' of step '{invocation.step.name}' is "
            f"{type(binding.value).__name__}, expected {value_type.__name__}."
        )
    return cast(T, binding.value)

def get_one_input[T](self, step: Step[Any] | str, value_type: type[T]) -> T:
    """Get the only recorded input of a runtime type before input coercion.

    All bound inputs are considered, including context values and defaults.
    Use `get_input` with a parameter name if more than one value matches.

    Args:
        step: A step from this pipeline or its name.
        value_type: Runtime type to match with `isinstance`.

    Returns:
        The single matching input value.

    Raises:
        KeyError: If the step is absent.
        LookupError: If zero or multiple inputs match `value_type`.
    """
    invocation = self.get_step_invocation(step)
    matches = [
        binding.value
        for binding in invocation.inputs.values()
        if isinstance(binding.value, value_type)
    ]
    if len(matches) != 1:
        raise LookupError(
            f"Expected exactly one input of type {value_type.__name__} "
            f"for step '{invocation.step.name}'; found {len(matches)}."
        )
    return cast(T, matches[0])