Fix tests

Author: dmontagu

pydantic_graph/pydantic_graph/beta/graph.py

@@ -22,7 +22,7 @@
 from pydantic_graph._utils import AbstractSpan, get_traceparent, logfire_span
 from pydantic_graph.beta.decision import Decision
 from pydantic_graph.beta.id_types import ForkStack, ForkStackItem, GraphRunId, JoinId, NodeId, NodeRunId, TaskId
-from pydantic_graph.beta.join import Join, Reducer
+from pydantic_graph.beta.join import Join, JoinNode, Reducer
 from pydantic_graph.beta.node import (
     EndNode,
     Fork,
@@ -441,7 +441,7 @@ def output(self) -> OutputT | None:
             return self._next.value
         return None
 
-    async def _iter_graph(
+    async def _iter_graph(  # noqa C901
         self,
     ) -> AsyncGenerator[
         EndMarker[OutputT] | JoinItem | Sequence[GraphTask], EndMarker[OutputT] | JoinItem | Sequence[GraphTask]
@@ -574,7 +574,7 @@ async def _handle_task(
                 step_context = StepContext[StateT, DepsT, Any](state, deps, inputs)
                 output = await node.call(step_context)
             if isinstance(node, NodeStep):
-                return self._handle_node(node, output, fork_stack)
+                return self._handle_node(output, fork_stack)
             else:
                 return self._handle_edges(node, output, fork_stack)
         elif isinstance(node, Join):
@@ -613,12 +613,13 @@ def _handle_decision(
 
     def _handle_node(
         self,
-        node_step: NodeStep[StateT, DepsT],
         next_node: BaseNode[StateT, DepsT, Any] | End[Any],
         fork_stack: ForkStack,
-    ) -> Sequence[GraphTask] | EndMarker[OutputT]:
+    ) -> Sequence[GraphTask] | JoinItem | EndMarker[OutputT]:
         if isinstance(next_node, StepNode):
             return [GraphTask(next_node.step.id, next_node.inputs, fork_stack)]
+        elif isinstance(next_node, JoinNode):
+            return JoinItem(next_node.join.id, next_node.inputs, fork_stack)
         elif isinstance(next_node, BaseNode):
             node_step = NodeStep(next_node.__class__)
             return [GraphTask(node_step.id, next_node, fork_stack)]
@@ -687,7 +688,10 @@ def _handle_path(self, path: Path, inputs: Any, fork_stack: ForkStack) -> Sequen
 
     def _handle_edges(self, node: AnyNode, inputs: Any, fork_stack: ForkStack) -> Sequence[GraphTask]:
         edges = self.graph.edges_by_source.get(node.id, [])
-        assert len(edges) == 1 or isinstance(node, Fork)  # this should have already been ensured during graph building
+        assert len(edges) == 1 or isinstance(node, Fork), (
+            edges,
+            node.id,
+        )  # this should have already been ensured during graph building
 
         new_tasks: list[GraphTask] = []
         for path in edges:

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -20,7 +20,7 @@
 from pydantic_graph.beta.decision import Decision, DecisionBranch, DecisionBranchBuilder
 from pydantic_graph.beta.graph import Graph
 from pydantic_graph.beta.id_types import ForkId, JoinId, NodeId
-from pydantic_graph.beta.join import Join, Reducer
+from pydantic_graph.beta.join import Join, JoinNode, Reducer
 from pydantic_graph.beta.node import (
     EndNode,
     Fork,
@@ -650,10 +650,21 @@ def _edge_from_return_hint(
                 )
                 if step is None:
                     raise exceptions.GraphSetupError(
-                        f'Node {node} return type hint includes a `StepNode` without a `Step` annotations. '
+                        f'Node {node} return type hint includes a `StepNode` without a `Step` annotation. '
                         'When returning `my_step.as_node()`, use `Annotated[StepNode[StateT, DepsT], my_step]` as the return type hint.'
                     )
                 destinations.append(step)
+            elif return_type_origin is JoinNode:
+                join = cast(
+                    Join[StateT, DepsT, Any, Any] | None,
+                    next((a for a in annotations if isinstance(a, Join)), None),  # pyright: ignore[reportUnknownArgumentType]
+                )
+                if join is None:
+                    raise exceptions.GraphSetupError(
+                        f'Node {node} return type hint includes a `JoinNode` without a `Join` annotation. '
+                        'When returning `my_join.as_node()`, use `Annotated[JoinNode[StateT, DepsT], my_join]` as the return type hint.'
+                    )
+                destinations.append(join)
             elif inspect.isclass(return_type_origin) and issubclass(return_type_origin, BaseNode):
                 destinations.append(NodeStep(return_type))
 

pydantic_graph/pydantic_graph/beta/join.py

@@ -9,10 +9,11 @@
 
 from abc import ABC
 from dataclasses import dataclass, field
-from typing import Generic
+from typing import Any, Generic, overload
 
 from typing_extensions import TypeVar
 
+from pydantic_graph import BaseNode, End, GraphRunContext
 from pydantic_graph.beta.id_types import ForkId, JoinId
 from pydantic_graph.beta.step import StepContext
 
@@ -259,3 +260,53 @@ def _force_covariant(self, inputs: InputT) -> OutputT:
             RuntimeError: Always raised as this method should never be called
         """
         raise RuntimeError('This method should never be called, it is just defined for typing purposes.')
+
+    @overload
+    def as_node(self, inputs: None = None) -> JoinNode[StateT, DepsT]: ...
+
+    @overload
+    def as_node(self, inputs: InputT) -> JoinNode[StateT, DepsT]: ...
+
+    def as_node(self, inputs: InputT | None = None) -> JoinNode[StateT, DepsT]:
+        """Create a step node with bound inputs.
+
+        Args:
+            inputs: The input data to bind to this step, or None
+
+        Returns:
+            A [`StepNode`][pydantic_graph.v2.step.StepNode] with this step and the bound inputs
+        """
+        return JoinNode(self, inputs)
+
+
+@dataclass
+class JoinNode(BaseNode[StateT, DepsT, Any]):
+    """A base node that represents a join item with bound inputs.
+
+    JoinNode bridges between the v1 and v2 graph execution systems by wrapping
+    a [`Join`][pydantic_graph.v2.step.Join] with bound inputs in a BaseNode interface.
+    It is not meant to be run directly but rather used to indicate transitions
+    to v2-style steps.
+    """
+
+    join: Join[StateT, DepsT, Any, Any]
+    """The step to execute."""
+
+    inputs: Any
+    """The inputs bound to this step."""
+
+    async def run(self, ctx: GraphRunContext[StateT, DepsT]) -> BaseNode[StateT, DepsT, Any] | End[Any]:
+        """Attempt to run the join node.
+
+        Args:
+            ctx: The graph execution context
+
+        Returns:
+            The result of step execution
+
+        Raises:
+            NotImplementedError: Always raised as StepNode is not meant to be run directly
+        """
+        raise NotImplementedError(
+            '`JoinNode` is not meant to be run directly, it is meant to be used in `BaseNode` subclasses to indicate a transition to v2-style steps.'
+        )

pydantic_graph/pydantic_graph/beta/step.py

@@ -76,6 +76,7 @@ def __repr__(self):
 
 
 if not TYPE_CHECKING:
+    # TODO: Try dropping inputs from StepContext, it would make for fewer generic params, could help
     StepContext = dataclass(StepContext)
 
 
@@ -186,6 +187,14 @@ def as_node(self, inputs: InputT | None = None) -> StepNode[StateT, DepsT]:
         """
         return StepNode(self, inputs)
 
+    def __repr__(self):
+        """Return a string representation of the step context.
+
+        Returns:
+            A string showing the class name and inputs
+        """
+        return f'Step(id={self.id!r}, call={self._call!r}, user_label={self.user_label!r})'
+
 
 @dataclass
 class StepNode(BaseNode[StateT, DepsT, Any]):

tests/graph/beta/test_decisions.py

@@ -27,12 +27,12 @@ async def choose_path(ctx: StepContext[DecisionState, None, None]) -> Literal['l
         return 'left'
 
     @g.step
-    async def left_path(ctx: StepContext[DecisionState, None, None]) -> str:
+    async def left_path(ctx: StepContext[DecisionState, None, object]) -> str:
         ctx.state.path_taken = 'left'
         return 'Went left'
 
     @g.step
-    async def right_path(ctx: StepContext[DecisionState, None, None]) -> str:
+    async def right_path(ctx: StepContext[DecisionState, None, object]) -> str:
         ctx.state.path_taken = 'right'
         return 'Went right'
 
@@ -258,11 +258,11 @@ async def choose(ctx: StepContext[DecisionState, None, None]) -> Literal['a', 'b
         return 'a'
 
     @g.step
-    async def path_a(ctx: StepContext[DecisionState, None, None]) -> str:
+    async def path_a(ctx: StepContext[DecisionState, None, object]) -> str:
         return 'Path A'
 
     @g.step
-    async def path_b(ctx: StepContext[DecisionState, None, None]) -> str:
+    async def path_b(ctx: StepContext[DecisionState, None, object]) -> str:
         return 'Path B'
 
     g.add(
@@ -285,15 +285,15 @@ async def test_decision_with_spread():
     g = GraphBuilder(state_type=DecisionState, output_type=int)
 
     @g.step
-    async def get_type(ctx: StepContext[DecisionState, None, None]) -> Literal['list', 'single']:
+    async def get_type(ctx: StepContext[DecisionState, None, object]) -> Literal['list', 'single']:
         return 'list'
 
     @g.step
-    async def make_list(ctx: StepContext[DecisionState, None, None]) -> list[int]:
+    async def make_list(ctx: StepContext[DecisionState, None, object]) -> list[int]:
         return [1, 2, 3]
 
     @g.step
-    async def make_single(ctx: StepContext[DecisionState, None, None]) -> int:
+    async def make_single(ctx: StepContext[DecisionState, None, object]) -> int:
         return 10
 
     @g.step
@@ -302,7 +302,7 @@ async def process_item(ctx: StepContext[DecisionState, None, int]) -> int:
         return ctx.inputs
 
     @g.step
-    async def get_value(ctx: StepContext[DecisionState, None, None]) -> int:
+    async def get_value(ctx: StepContext[DecisionState, None, object]) -> int:
         return ctx.state.value
 
     g.add(
@@ -321,4 +321,5 @@ async def get_value(ctx: StepContext[DecisionState, None, None]) -> int:
     graph = g.build()
     state = DecisionState()
     result = await graph.run(state=state)
+    assert result == 6
     assert state.value == 6  # 1 + 2 + 3

tests/graph/beta/test_edge_cases.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Any
 
 import pytest
 
@@ -30,7 +31,7 @@ async def test_graph_with_no_steps():
 
 async def test_step_returning_none():
     """Test steps that return None."""
-    g = GraphBuilder(state_type=EdgeCaseState, output_type=None)
+    g = GraphBuilder(state_type=EdgeCaseState)
 
     @g.step
     async def do_nothing(ctx: StepContext[EdgeCaseState, None, None]) -> None:
@@ -176,7 +177,7 @@ async def test_long_sequential_chain():
     """Test a long chain of sequential steps."""
     g = GraphBuilder(state_type=EdgeCaseState, output_type=int)
 
-    steps = []
+    steps: list[Any] = []
     for i in range(10):
 
         @g.step(node_id=f'step_{i}')
@@ -223,7 +224,7 @@ async def single_source(ctx: StepContext[EdgeCaseState, None, None]) -> int:
 
 async def test_null_reducer_with_no_inputs():
     """Test NullReducer behavior with spread that produces no items."""
-    g = GraphBuilder(state_type=EdgeCaseState, output_type=None)
+    g = GraphBuilder(state_type=EdgeCaseState)
 
     @g.step
     async def empty_list(ctx: StepContext[EdgeCaseState, None, None]) -> list[int]:
@@ -315,11 +316,7 @@ async def test_state_with_mutable_collections():
 
     @dataclass
     class MutableState:
-        items: list[int] = None  # type: ignore
-
-        def __post_init__(self):
-            if self.items is None:
-                self.items = []
+        items: list[int] = field(default_factory=list)
 
     g = GraphBuilder(state_type=MutableState, output_type=list[int])
 

tests/graph/beta/test_edge_labels.py

@@ -132,15 +132,15 @@ async def test_label_on_decision_branch():
     g = GraphBuilder(state_type=LabelState, output_type=str)
 
     @g.step
-    async def choose(ctx: StepContext[LabelState, None, None]) -> Literal['a', 'b']:
+    async def choose(ctx: StepContext[LabelState, None, object]) -> Literal['a', 'b']:
         return 'a'
 
     @g.step
-    async def path_a(ctx: StepContext[LabelState, None, None]) -> str:
+    async def path_a(ctx: StepContext[LabelState, None, object]) -> str:
         return 'A'
 
     @g.step
-    async def path_b(ctx: StepContext[LabelState, None, None]) -> str:
+    async def path_b(ctx: StepContext[LabelState, None, object]) -> str:
         return 'B'
 
     g.add(

tests/graph/beta/test_graph_builder.py

@@ -178,7 +178,7 @@ async def multiply(ctx: StepContext[SimpleState, MyDeps, None]) -> int:
 
 async def test_empty_graph():
     """Test that a minimal graph can be built and run."""
-    g = GraphBuilder(output_type=int)
+    g = GraphBuilder(input_type=int, output_type=int)
 
     g.add(g.edge_from(g.start_node).to(g.end_node))
 
@@ -208,7 +208,7 @@ async def return_value(ctx: StepContext[None, None, None]) -> int:
 
 async def test_explicit_graph_name():
     """Test setting an explicit graph name."""
-    g = GraphBuilder(name='ExplicitName', output_type=int)
+    g = GraphBuilder(name='ExplicitName', input_type=int, output_type=int)
 
     g.add(g.edge_from(g.start_node).to(g.end_node))
 

tests/graph/beta/test_graph_iteration.py

@@ -3,11 +3,13 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Any
 
 import pytest
 
 from pydantic_graph.beta import GraphBuilder, StepContext
 from pydantic_graph.beta.graph import EndMarker, GraphTask, JoinItem
+from pydantic_graph.beta.id_types import NodeId
 
 pytestmark = pytest.mark.anyio
 
@@ -39,14 +41,15 @@ async def double(ctx: StepContext[IterState, None, int]) -> int:
     graph = g.build()
     state = IterState()
 
-    events = []
+    events: list[Any] = []
     async with graph.iter(state=state) as run:
         async for event in run:
             events.append(event)
 
     assert len(events) > 0
-    assert isinstance(events[-1], EndMarker)
-    assert events[-1].value == 2
+    last_event = events[-1]
+    assert isinstance(last_event, EndMarker)
+    assert last_event.value == 2  # pyright: ignore[reportUnknownMemberType]
 
 
 async def test_iter_with_next():
@@ -99,7 +102,7 @@ async def my_step(ctx: StepContext[IterState, None, None]) -> int:
     graph = g.build()
     state = IterState()
 
-    task_nodes = []
+    task_nodes: list[NodeId] = []
     async with graph.iter(state=state) as run:
         async for event in run:
             if isinstance(event, list):
@@ -276,7 +279,7 @@ async def step_three(ctx: StepContext[IterState, None, int]) -> int:
 
     async with graph.iter(state=state) as run:
         event_count = 0
-        async for event in run:
+        async for _ in run:
             event_count += 1
             if event_count >= 2:
                 break  # Early termination
@@ -308,9 +311,9 @@ async def double_counter(ctx: StepContext[IterState, None, None]) -> int:
     graph = g.build()
     state = IterState()
 
-    state_snapshots = []
+    state_snapshots: list[Any] = []
     async with graph.iter(state=state) as run:
-        async for event in run:
+        async for _ in run:
             # Take a snapshot of the state after each event
             state_snapshots.append(state.counter)
 

tests/graph/beta/test_joins_and_reducers.py

@@ -18,7 +18,7 @@ class SimpleState:
 
 async def test_null_reducer():
     """Test NullReducer that discards all inputs."""
-    g = GraphBuilder(state_type=SimpleState, output_type=None)
+    g = GraphBuilder(state_type=SimpleState)
 
     @g.step
     async def source(ctx: StepContext[SimpleState, None, None]) -> list[int]: