Make StepContext a dataclass again

Author: dmontagu

pydantic_graph/pydantic_graph/beta/graph.py

@@ -492,7 +492,7 @@ def _handle_result(result: EndMarker[OutputT] | JoinItem | Sequence[GraphTask])
                     reducer, _ = reducer_and_fork_stack
 
                 try:
-                    reducer.reduce(StepContext(self.state, self.deps, result.inputs))
+                    reducer.reduce(StepContext(state=self.state, deps=self.deps, inputs=result.inputs))
                 except StopIteration:
                     # cancel all concurrently running tasks with the same fork_run_id of the parent fork
                     task_ids_to_cancel = set[TaskID]()
@@ -522,7 +522,7 @@ def _handle_result(result: EndMarker[OutputT] | JoinItem | Sequence[GraphTask])
 
                     for join_id, fork_run_id in self._get_completed_fork_runs(source_task, tasks_by_id.values()):
                         reducer, fork_stack = self._active_reducers.pop((join_id, fork_run_id))
-                        output = reducer.finalize(StepContext(self.state, self.deps, None))
+                        output = reducer.finalize(StepContext(state=self.state, deps=self.deps, inputs=None))
                         join_node = self.graph.nodes[join_id]
                         assert isinstance(
                             join_node, Join
@@ -545,7 +545,7 @@ def _handle_result(result: EndMarker[OutputT] | JoinItem | Sequence[GraphTask])
                         continue  # this reducer is a strict prefix for one of the other active reducers
 
                     self._active_reducers.pop((join_id, fork_run_id))  # we're finalizing it now
-                    output = reducer.finalize(StepContext(self.state, self.deps, None))
+                    output = reducer.finalize(StepContext(state=self.state, deps=self.deps, inputs=None))
                     join_node = self.graph.nodes[join_id]
                     assert isinstance(join_node, Join)  # We could drop this but if it fails it means there is a bug.
                     new_tasks = self._handle_edges(join_node, output, fork_stack)
@@ -576,7 +576,7 @@ async def _handle_task(
                 if self.graph.auto_instrument:
                     stack.enter_context(logfire_span('run node {node_id}', node_id=node.id, node=node))
 
-                step_context = StepContext[StateT, DepsT, Any](state, deps, inputs)
+                step_context = StepContext[StateT, DepsT, Any](state=state, deps=deps, inputs=inputs)
                 output = await node.call(step_context)
             if isinstance(node, NodeStep):
                 return self._handle_node(output, fork_stack)
@@ -684,7 +684,7 @@ def _handle_path(self, path: Path, inputs: Any, fork_stack: ForkStack) -> Sequen
         elif isinstance(item, BroadcastMarker):
             return [GraphTask(item.fork_id, inputs, fork_stack)]
         elif isinstance(item, TransformMarker):
-            inputs = item.transform(StepContext(self.state, self.deps, inputs))
+            inputs = item.transform(StepContext(state=self.state, deps=self.deps, inputs=inputs))
             return self._handle_path(path.next_path, inputs, fork_stack)
         elif isinstance(item, LabelMarker):
             return self._handle_path(path.next_path, inputs, fork_stack)

pydantic_graph/pydantic_graph/beta/step.py

@@ -8,8 +8,8 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Generic, Protocol, cast, get_origin, overload
+from dataclasses import dataclass, field
+from typing import Any, Generic, Protocol, cast, get_origin, overload
 
 from typing_extensions import TypeVar
 
@@ -22,6 +22,7 @@
 OutputT = TypeVar('OutputT', infer_variance=True)
 
 
+@dataclass(kw_only=True, frozen=True)
 class StepContext(Generic[StateT, DepsT, InputT]):
     """Context information passed to step functions during graph execution.
 
@@ -35,49 +36,14 @@ class StepContext(Generic[StateT, DepsT, InputT]):
         InputT: The type of the input data
     """
 
-    if TYPE_CHECKING:
+    state: StateT = field(repr=False)  # exclude from repr to keep things concise
+    """The current graph state."""
 
-        def __init__(self, state: StateT, deps: DepsT, inputs: InputT):
-            self._state = state
-            self._deps = deps
-            self._inputs = inputs
+    deps: DepsT = field(repr=False)  # exclude from repr to keep things concise
+    """The dependencies available to this step."""
 
-        @property
-        def state(self) -> StateT:
-            """The current graph state."""
-            return self._state
-
-        @property
-        def deps(self) -> DepsT:
-            """The dependencies available to this step."""
-            return self._deps
-
-        @property
-        def inputs(self) -> InputT:
-            """The input data for this step."""
-            return self._inputs
-    else:
-        state: StateT
-        """The current graph state."""
-
-        deps: DepsT
-        """The dependencies available to this step."""
-
-        inputs: InputT
-        """The input data for this step."""
-
-    def __repr__(self) -> str:
-        """Return a string representation of the step context.
-
-        Returns:
-            A string showing the class name and inputs
-        """
-        return f'{self.__class__.__name__}(inputs={self.inputs})'
-
-
-if not TYPE_CHECKING:
-    # TODO: Try dropping inputs from StepContext, it would make for fewer generic params, could help
-    StepContext = dataclass(StepContext)
+    inputs: InputT
+    """The input data for this step."""
 
 
 class StepFunction(Protocol[StateT, DepsT, InputT, OutputT]):