WIP

Author: dmontagu

pydantic_graph/pydantic_graph/beta/decision.py

@@ -9,13 +9,12 @@
 
 from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Generic
+from typing import TYPE_CHECKING, Any, Final, Generic
 
 from typing_extensions import Never, Self, TypeVar
 
 from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID
-from pydantic_graph.beta.paths import Path, PathBuilder
-from pydantic_graph.beta.step import StepFunction
+from pydantic_graph.beta.paths import Path, PathBuilder, TransformFunction
 from pydantic_graph.beta.util import TypeOrTypeExpression
 
 if TYPE_CHECKING:
@@ -124,24 +123,26 @@ class DecisionBranch(Generic[SourceT]):
 """Type variable for transformed output."""
 
 
-@dataclass
+@dataclass(kw_only=True)
 class DecisionBranchBuilder(Generic[StateT, DepsT, OutputT, SourceT, HandledT]):
     """Builder for constructing decision branches with fluent API.
 
     This builder provides methods to configure branches with destinations,
     forks, and transformations in a type-safe manner.
     """
 
-    decision: Decision[StateT, DepsT, HandledT]
+    # The use of `Final` on these attributes is necessary for them to be treated as read-only for purposes
+    # of variance-inference. This could be done with `frozen` but that
+    decision: Final[Decision[StateT, DepsT, HandledT]]
     """The parent decision node."""
 
-    source: TypeOrTypeExpression[SourceT]
+    source: Final[TypeOrTypeExpression[SourceT]]
     """The expected source type for this branch."""
 
-    matches: Callable[[Any], bool] | None
+    matches: Final[Callable[[Any], bool] | None]
     """Optional matching predicate."""
 
-    path_builder: PathBuilder[StateT, DepsT, OutputT]
+    path_builder: Final[PathBuilder[StateT, DepsT, OutputT]]
     """Builder for the execution path."""
 
     @property
@@ -194,7 +195,7 @@ def fork(
         return DecisionBranch(source=self.source, matches=self.matches, path=self.path_builder.fork(new_paths))
 
     def transform(
-        self, func: StepFunction[StateT, DepsT, OutputT, NewOutputT], /
+        self, func: TransformFunction[StateT, DepsT, OutputT, NewOutputT], /
     ) -> DecisionBranchBuilder[StateT, DepsT, NewOutputT, SourceT, HandledT]:
         """Apply a transformation to the branch's output.
 

pydantic_graph/pydantic_graph/beta/paths.py

@@ -13,20 +13,49 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Generic, get_origin, overload
 
-from typing_extensions import Self, TypeAliasType, TypeVar
+from typing_extensions import Protocol, Self, TypeAliasType, TypeVar
 
 from pydantic_graph import BaseNode
 from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID
-from pydantic_graph.beta.step import NodeStep, StepFunction
+from pydantic_graph.beta.step import NodeStep, StepContext
 
 StateT = TypeVar('StateT', infer_variance=True)
 DepsT = TypeVar('DepsT', infer_variance=True)
 OutputT = TypeVar('OutputT', infer_variance=True)
+InputT = TypeVar('InputT', infer_variance=True)
 
 if TYPE_CHECKING:
     from pydantic_graph.beta.node_types import AnyDestinationNode, DestinationNode, SourceNode
 
 
+class TransformFunction(Protocol[StateT, DepsT, InputT, OutputT]):
+    """Protocol for step functions that can be executed in the graph.
+
+    Transform functions are sync callables that receive a step context and return
+    a result. This protocol enables serialization and deserialization of step
+    calls similar to how evaluators work.
+
+    This is very similar to a StepFunction, but must be sync instead of async.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        InputT: The type of the input data
+        OutputT: The type of the output data
+    """
+
+    def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> OutputT:
+        """Execute the step function with the given context.
+
+        Args:
+            ctx: The step context containing state, dependencies, and inputs
+
+        Returns:
+            An awaitable that resolves to the step's output
+        """
+        raise NotImplementedError
+
+
 @dataclass
 class TransformMarker:
     """A marker indicating a data transformation step in a path.
@@ -35,7 +64,7 @@ class TransformMarker:
     through the graph path.
     """
 
-    transform: StepFunction[Any, Any, Any, Any]
+    transform: TransformFunction[Any, Any, Any, Any]
     """The step function that performs the transformation."""
 
 
@@ -196,7 +225,7 @@ def fork(self, forks: Sequence[Path], /, *, fork_id: str | None = None) -> Path:
         next_item = BroadcastMarker(paths=forks, fork_id=ForkID(NodeID(fork_id or 'broadcast_' + secrets.token_hex(8))))
         return Path(items=[*self.working_items, next_item])
 
-    def transform(self, func: StepFunction[StateT, DepsT, OutputT, Any], /) -> PathBuilder[StateT, DepsT, Any]:
+    def transform(self, func: TransformFunction[StateT, DepsT, OutputT, Any], /) -> PathBuilder[StateT, DepsT, Any]:
         """Add a transformation step to the path.
 
         Args:
@@ -385,7 +414,7 @@ def map(
             path_builder=self.path_builder.map(fork_id=fork_id, downstream_join_id=downstream_join_id),
         )
 
-    def transform(self, func: StepFunction[StateT, DepsT, OutputT, Any], /) -> EdgePathBuilder[StateT, DepsT, Any]:
+    def transform(self, func: TransformFunction[StateT, DepsT, OutputT, Any], /) -> EdgePathBuilder[StateT, DepsT, Any]:
         """Add a transformation step to the edge path.
 
         Args:

tests/graph/beta/test_decisions.py

@@ -3,11 +3,11 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Literal
+from typing import Any, Literal
 
 import pytest
 
-from pydantic_graph.beta import GraphBuilder, Reducer, StepContext, TypeExpression
+from pydantic_graph.beta import GraphBuilder, ListReducer, Reducer, StepContext, TypeExpression
 
 pytestmark = pytest.mark.anyio
 
@@ -350,27 +350,30 @@ async def return_list(ctx: StepContext[DecisionState, None, None]) -> list[int]:
     async def process_item(ctx: StepContext[DecisionState, None, int]) -> int:
         return ctx.inputs * 2
 
-    class SumReducer(Reducer[object, object, float, float]):
+    class SumReducer(Reducer[object, object, int, int]):
         """A reducer that sums values."""
 
-        value: float = 0.0
+        value: int = 0
 
-        def reduce(self, ctx: StepContext[object, object, float]) -> None:
+        def reduce(self, ctx: StepContext[object, object, int]) -> None:
             self.value += ctx.inputs
 
-        def finalize(self, ctx: StepContext[object, object, None]) -> float:
+        def finalize(self, ctx: StepContext[object, object, None]) -> int:
             return self.value
 
     sum_results = g.join(SumReducer)
 
+    def is_list_int(x: Any) -> bool:
+        return isinstance(x, list) and all(isinstance(y, int) for y in x)  # pyright: ignore[reportUnknownVariableType]
+
     # Use decision with map to test last_fork_id
     g.add(
         g.edge_from(g.start_node).to(return_list),
         g.edge_from(return_list).to(
             g.decision().branch(
                 g.match(
                     TypeExpression[list[int]],
-                    matches=lambda x: isinstance(x, list) and all(isinstance(y, int) for y in x),
+                    matches=is_list_int,
                 )
                 .map()
                 .to(process_item)
@@ -397,8 +400,8 @@ async def get_value(ctx: StepContext[DecisionState, None, None]) -> int:
     async def format_result(ctx: StepContext[DecisionState, None, str]) -> str:
         return f'Result: {ctx.inputs}'
 
-    async def double_value(ctx: StepContext[DecisionState, None, int], value: int) -> str:
-        return str(value * 2)
+    def double_value(ctx: StepContext[DecisionState, None, int]) -> str:
+        return str(ctx.inputs * 2)
 
     g.add(
         g.edge_from(g.start_node).to(get_value),
@@ -458,6 +461,8 @@ async def path_1(ctx: StepContext[DecisionState, None, object]) -> str:
     async def path_2(ctx: StepContext[DecisionState, None, object]) -> str:
         return 'Path 2'
 
+    collect = g.join(ListReducer[str])
+
     @g.step
     async def combine(ctx: StepContext[DecisionState, None, list[str]]) -> str:
         return ', '.join(ctx.inputs)
@@ -474,7 +479,8 @@ async def combine(ctx: StepContext[DecisionState, None, list[str]]) -> str:
                 )
             )
         ),
-        g.edge_from(path_1, path_2).join().to(combine),
+        g.edge_from(path_1, path_2).to(collect),
+        g.edge_from(collect).to(combine),
         g.edge_from(combine).to(g.end_node),
     )