Fix a bug

Author: dmontagu

docs/graph/beta/parallel.md

@@ -114,7 +114,7 @@ The convenience method [`add_mapping_edge()`][pydantic_graph.beta.graph_builder.
 ```python {title="mapping_convenience.py"}
 from dataclasses import dataclass
 
-from pydantic_graph.beta import GraphBuilder, ListReducer, StepContext
+from pydantic_graph.beta import GraphBuilder, ListReducer, StepContext, Reducer
 
 
 @dataclass

pydantic_graph/pydantic_graph/beta/decision.py

@@ -65,7 +65,7 @@ def branch(self, branch: DecisionBranch[T]) -> Decision[StateT, DepsT, HandledT
         """
         return Decision(id=self.id, branches=self.branches + [branch], note=self.note)
 
-    def _force_handled_contravariant(self, inputs: HandledT) -> Never:
+    def _force_handled_contravariant(self, inputs: HandledT) -> Never:  # pragma: no cover
         """Forces this type to be contravariant in the HandledT type variable.
 
         This is an implementation detail of how we can type-check that all possible input types have

pydantic_graph/pydantic_graph/beta/graph.py

@@ -34,8 +34,8 @@
     BroadcastMarker,
     DestinationMarker,
     LabelMarker,
+    MapMarker,
     Path,
-    SpreadMarker,
     TransformMarker,
 )
 from pydantic_graph.beta.step import NodeStep, Step, StepContext, StepNode
@@ -654,7 +654,7 @@ def _handle_path(self, path: Path, inputs: Any, fork_stack: ForkStack) -> Sequen
         item = path.items[0]
         if isinstance(item, DestinationMarker):
             return [GraphTask(item.destination_id, inputs, fork_stack)]
-        elif isinstance(item, SpreadMarker):
+        elif isinstance(item, MapMarker):
             # Eagerly raise a clear error if the input value is not iterable as expected
             try:
                 iter(inputs)

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -38,9 +38,9 @@
     DestinationMarker,
     EdgePath,
     EdgePathBuilder,
+    MapMarker,
     Path,
     PathBuilder,
-    SpreadMarker,
 )
 from pydantic_graph.beta.step import NodeStep, Step, StepFunction, StepNode
 from pydantic_graph.beta.util import TypeOrTypeExpression, get_callable_name, unpack_type_expression
@@ -362,7 +362,7 @@ def _handle_path(p: Path):
                     self._insert_node(new_node)
                     for path in item.paths:
                         _handle_path(Path(items=[*path.items]))
-                elif isinstance(item, SpreadMarker):
+                elif isinstance(item, MapMarker):
                     new_node = Fork[Any, Any](id=item.fork_id, is_map=True)
                     self._insert_node(new_node)
                 elif isinstance(item, DestinationMarker):
@@ -376,6 +376,9 @@ def _handle_path(p: Path):
             for destination_node in edge.destinations:
                 destinations.append(destination_node)
                 self._insert_node(destination_node)
+                if isinstance(destination_node, Decision):
+                    for branch in destination_node.branches:
+                        _handle_path(branch.path)
 
             _handle_path(edge.path)
 
@@ -570,6 +573,7 @@ def _get_new_decision_id(self) -> str:
             self._decision_index += 1
         return node_id
 
+    # TODO(P1): Need to use or remove this..
     def _get_new_broadcast_id(self, from_: str | None = None) -> str:
         """Generate a unique ID for a new broadcast fork.
 
@@ -590,6 +594,7 @@ def _get_new_broadcast_id(self, from_: str | None = None) -> str:
             index += 1
         return node_id
 
+    # TODO(P1): Need to use or remove this..
     def _get_new_map_id(self, from_: str | None = None, to: str | None = None) -> str:
         """Generate a unique ID for a new map fork.
 
@@ -758,7 +763,7 @@ def _normalize_forks(
     while paths_to_handle:
         path = paths_to_handle.pop()
         for item in path.items:
-            if isinstance(item, SpreadMarker):
+            if isinstance(item, MapMarker):
                 assert item.fork_id in new_nodes
                 new_edges[item.fork_id] = [path.next_path]
             if isinstance(item, BroadcastMarker):
@@ -811,7 +816,7 @@ def _handle_path(path: Path, last_source_id: NodeID):
                 last_source_id: The current source node ID
             """
             for item in path.items:
-                if isinstance(item, SpreadMarker):
+                if isinstance(item, MapMarker):
                     fork_ids.add(item.fork_id)
                     edges[last_source_id].append(item.fork_id)
                     last_source_id = item.fork_id

pydantic_graph/pydantic_graph/beta/join.py

@@ -242,7 +242,7 @@ def create_reducer(self) -> Reducer[StateT, DepsT, InputT, OutputT]:
     # def deserialize_reducer(self, serialized: bytes) -> Reducer[InputT, OutputT]:
     #     return self._type_adapter.validate_json(serialized)
 
-    def _force_covariant(self, inputs: InputT) -> OutputT:
+    def _force_covariant(self, inputs: InputT) -> OutputT:  # pragma: no cover
         """Force covariant typing for generic parameters.
 
         This method exists solely for typing purposes and should never be called.

pydantic_graph/pydantic_graph/beta/mermaid.py

@@ -11,7 +11,7 @@
 from pydantic_graph.beta.id_types import NodeID
 from pydantic_graph.beta.join import Join
 from pydantic_graph.beta.node import EndNode, Fork, StartNode
-from pydantic_graph.beta.paths import BroadcastMarker, DestinationMarker, LabelMarker, Path, SpreadMarker
+from pydantic_graph.beta.paths import BroadcastMarker, DestinationMarker, LabelMarker, MapMarker, Path
 from pydantic_graph.beta.step import NodeStep, Step
 
 DEFAULT_HIGHLIGHT_CSS = 'fill:#fdff32'
@@ -57,7 +57,7 @@ def build_mermaid_graph(graph: Graph[Any, Any, Any, Any]) -> MermaidGraph:  # no
     def _collect_edges(path: Path, last_source_id: NodeID) -> None:
         working_label: str | None = None
         for item in path.items:
-            if isinstance(item, SpreadMarker):
+            if isinstance(item, MapMarker):
                 edges_by_source[last_source_id].append(MermaidEdge(last_source_id, item.fork_id, working_label))
                 return  # map markers correspond to nodes already in the graph; downstream gets handled separately
             elif isinstance(item, BroadcastMarker):

pydantic_graph/pydantic_graph/beta/node.py

@@ -44,7 +44,7 @@ class EndNode(Generic[InputT]):
     id = NodeID('__end__')
     """Fixed identifier for the end node."""
 
-    def _force_variance(self, inputs: InputT) -> None:
+    def _force_variance(self, inputs: InputT) -> None:  # pragma: no cover
         """Force type variance for proper generic typing.
 
         This method exists solely for type checking purposes and should never be called.
@@ -57,9 +57,6 @@ def _force_variance(self, inputs: InputT) -> None:
         """
         raise RuntimeError('This method should never be called, it is just defined for typing purposes.')
 
-    # def _force_variance(self) -> InputT:
-    #     raise RuntimeError('This method should never be called, it is just defined for typing purposes.')
-
 
 @dataclass
 class Fork(Generic[InputT, OutputT]):
@@ -80,7 +77,7 @@ class Fork(Generic[InputT, OutputT]):
     If False, InputT must be OutputT and the same data is sent to all branches.
     """
 
-    def _force_variance(self, inputs: InputT) -> OutputT:
+    def _force_variance(self, inputs: InputT) -> OutputT:  # pragma: no cover
         """Force type variance for proper generic typing.
 
         This method exists solely for type checking purposes and should never be called.

pydantic_graph/pydantic_graph/beta/paths.py

@@ -40,7 +40,7 @@ class TransformMarker:
 
 
 @dataclass
-class SpreadMarker:
+class MapMarker:
     """A marker indicating that iterable data should be map across parallel paths.
 
     Spread markers take iterable input and create parallel execution paths
@@ -92,7 +92,7 @@ class DestinationMarker:
     """The unique identifier of the destination node."""
 
 
-PathItem = TypeAliasType('PathItem', TransformMarker | SpreadMarker | BroadcastMarker | LabelMarker | DestinationMarker)
+PathItem = TypeAliasType('PathItem', TransformMarker | MapMarker | BroadcastMarker | LabelMarker | DestinationMarker)
 """Type alias for any item that can appear in a path sequence."""
 
 
@@ -108,14 +108,14 @@ class Path:
     """The sequence of path items that define this path."""
 
     @property
-    def last_fork(self) -> BroadcastMarker | SpreadMarker | None:
+    def last_fork(self) -> BroadcastMarker | MapMarker | None:
         """Get the most recent fork or map marker in this path.
 
         Returns:
-            The last BroadcastMarker or SpreadMarker in the path, or None if no forks exist
+            The last BroadcastMarker or MapMarker in the path, or None if no forks exist
         """
         for item in reversed(self.items):
-            if isinstance(item, BroadcastMarker | SpreadMarker):
+            if isinstance(item, BroadcastMarker | MapMarker):
                 return item
         return None
 
@@ -146,14 +146,14 @@ class PathBuilder(Generic[StateT, DepsT, OutputT]):
     """The accumulated sequence of path items being built."""
 
     @property
-    def last_fork(self) -> BroadcastMarker | SpreadMarker | None:
+    def last_fork(self) -> BroadcastMarker | MapMarker | None:
         """Get the most recent fork or map marker in the working path.
 
         Returns:
-            The last BroadcastMarker or SpreadMarker in the working items, or None if no forks exist
+            The last BroadcastMarker or MapMarker in the working items, or None if no forks exist
         """
         for item in reversed(self.working_items):
-            if isinstance(item, BroadcastMarker | SpreadMarker):
+            if isinstance(item, BroadcastMarker | MapMarker):
                 return item
         return None
 
@@ -226,7 +226,7 @@ def map(
         Returns:
             A new PathBuilder that operates on individual items from the iterable
         """
-        next_item = SpreadMarker(
+        next_item = MapMarker(
             fork_id=NodeID(fork_id or 'map_' + secrets.token_hex(8)), downstream_join_id=downstream_join_id
         )
         return PathBuilder[StateT, DepsT, Any](working_items=[*self.working_items, next_item])

tests/graph/beta/test_decisions.py

@@ -7,7 +7,7 @@
 
 import pytest
 
-from pydantic_graph.beta import GraphBuilder, StepContext, TypeExpression
+from pydantic_graph.beta import GraphBuilder, Reducer, StepContext, TypeExpression
 
 pytestmark = pytest.mark.anyio
 
@@ -323,3 +323,162 @@ async def get_value(ctx: StepContext[DecisionState, None, object]) -> int:
     result = await graph.run(state=state)
     assert result == 6
     assert state.value == 6  # 1 + 2 + 3
+
+
+async def test_decision_branch_last_fork_id_none():
+    """Test DecisionBranchBuilder.last_fork_id when there are no forks."""
+    from pydantic_graph.beta.decision import Decision, DecisionBranchBuilder
+    from pydantic_graph.beta.id_types import NodeID
+    from pydantic_graph.beta.paths import PathBuilder
+
+    decision = Decision[DecisionState, None, int](id=NodeID('test'), branches=[], note=None)
+    path_builder = PathBuilder[DecisionState, None, int](working_items=[])
+    branch_builder = DecisionBranchBuilder(decision=decision, source=int, matches=None, path_builder=path_builder)
+
+    assert branch_builder.last_fork_id is None
+
+
+async def test_decision_branch_last_fork_id_with_map():
+    """Test DecisionBranchBuilder.last_fork_id after a map operation."""
+    g = GraphBuilder(state_type=DecisionState, output_type=int)
+
+    @g.step
+    async def return_list(ctx: StepContext[DecisionState, None, None]) -> list[int]:
+        return [1, 2, 3]
+
+    @g.step
+    async def process_item(ctx: StepContext[DecisionState, None, int]) -> int:
+        return ctx.inputs * 2
+
+    class SumReducer(Reducer[object, object, float, float]):
+        """A reducer that sums values."""
+
+        value: float = 0.0
+
+        def reduce(self, ctx: StepContext[object, object, float]) -> None:
+            self.value += ctx.inputs
+
+        def finalize(self, ctx: StepContext[object, object, None]) -> float:
+            return self.value
+
+    sum_results = g.join(SumReducer)
+
+    # 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),
+                )
+                .map()
+                .to(process_item)
+            )
+        ),
+        g.edge_from(process_item).to(sum_results),
+        g.edge_from(sum_results).to(g.end_node),
+    )
+
+    graph = g.build()
+    result = await graph.run(state=DecisionState())
+    assert result == 12  # (1+2+3) * 2
+
+
+async def test_decision_branch_transform():
+    """Test DecisionBranchBuilder.transform method."""
+    g = GraphBuilder(state_type=DecisionState, output_type=str)
+
+    @g.step
+    async def get_value(ctx: StepContext[DecisionState, None, None]) -> int:
+        return 10
+
+    @g.step
+    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)
+
+    g.add(
+        g.edge_from(g.start_node).to(get_value),
+        g.edge_from(get_value).to(g.decision().branch(g.match(int).transform(double_value).to(format_result))),
+        g.edge_from(format_result).to(g.end_node),
+    )
+
+    graph = g.build()
+    result = await graph.run(state=DecisionState())
+    assert result == 'Result: 20'
+
+
+async def test_decision_branch_label():
+    """Test DecisionBranchBuilder.label method."""
+    g = GraphBuilder(state_type=DecisionState, output_type=str)
+
+    @g.step
+    async def get_value(ctx: StepContext[DecisionState, None, None]) -> Literal['a', 'b']:
+        return 'a'
+
+    @g.step
+    async def handle_a(ctx: StepContext[DecisionState, None, object]) -> str:
+        return 'Got A'
+
+    @g.step
+    async def handle_b(ctx: StepContext[DecisionState, None, object]) -> str:
+        return 'Got B'
+
+    g.add(
+        g.edge_from(g.start_node).to(get_value),
+        g.edge_from(get_value).to(
+            g.decision()
+            .branch(g.match(TypeExpression[Literal['a']]).label('path A').to(handle_a))
+            .branch(g.match(TypeExpression[Literal['b']]).label('path B').to(handle_b))
+        ),
+        g.edge_from(handle_a, handle_b).to(g.end_node),
+    )
+
+    graph = g.build()
+    result = await graph.run(state=DecisionState())
+    assert result == 'Got A'
+
+
+async def test_decision_branch_fork():
+    """Test DecisionBranchBuilder.fork method."""
+    g = GraphBuilder(state_type=DecisionState, output_type=str)
+
+    @g.step
+    async def choose_option(ctx: StepContext[DecisionState, None, None]) -> Literal['fork']:
+        return 'fork'
+
+    @g.step
+    async def path_1(ctx: StepContext[DecisionState, None, object]) -> str:
+        return 'Path 1'
+
+    @g.step
+    async def path_2(ctx: StepContext[DecisionState, None, object]) -> str:
+        return 'Path 2'
+
+    @g.step
+    async def combine(ctx: StepContext[DecisionState, None, list[str]]) -> str:
+        return ', '.join(ctx.inputs)
+
+    g.add(
+        g.edge_from(g.start_node).to(choose_option),
+        g.edge_from(choose_option).to(
+            g.decision().branch(
+                g.match(TypeExpression[Literal['fork']]).fork(
+                    lambda b: [
+                        b.decision.branch(g.match(TypeExpression[Literal['fork']]).to(path_1)),
+                        b.decision.branch(g.match(TypeExpression[Literal['fork']]).to(path_2)),
+                    ]
+                )
+            )
+        ),
+        g.edge_from(path_1, path_2).join().to(combine),
+        g.edge_from(combine).to(g.end_node),
+    )
+
+    graph = g.build()
+    result = await graph.run(state=DecisionState())
+    assert 'Path 1' in result
+    assert 'Path 2' in result