Fix failing docs tests

Author: dmontagu

docs/graph/beta/joins.md

@@ -25,28 +25,29 @@ class SimpleState:
     pass
 
 
-async def main():
-    g = GraphBuilder(state_type=SimpleState, output_type=list[int])
+g = GraphBuilder(state_type=SimpleState, output_type=list[int])
 
-    @g.step
-    async def generate_numbers(ctx: StepContext[SimpleState, None, None]) -> list[int]:
-        return [1, 2, 3, 4, 5]
+@g.step
+async def generate_numbers(ctx: StepContext[SimpleState, None, None]) -> list[int]:
+    return [1, 2, 3, 4, 5]
 
-    @g.step
-    async def square(ctx: StepContext[SimpleState, None, int]) -> int:
-        return ctx.inputs * ctx.inputs
+@g.step
+async def square(ctx: StepContext[SimpleState, None, int]) -> int:
+    return ctx.inputs * ctx.inputs
 
-    # Create a join to collect all squared values
-    collect = g.join(ListReducer[int])
+# Create a join to collect all squared values
+collect = g.join(ListReducer[int])
 
-    g.add(
-        g.edge_from(g.start_node).to(generate_numbers),
-        g.edge_from(generate_numbers).map().to(square),
-        g.edge_from(square).to(collect),
-        g.edge_from(collect).to(g.end_node),
-    )
+g.add(
+    g.edge_from(g.start_node).to(generate_numbers),
+    g.edge_from(generate_numbers).map().to(square),
+    g.edge_from(square).to(collect),
+    g.edge_from(collect).to(g.end_node),
+)
 
-    graph = g.build()
+graph = g.build()
+
+async def main():
     result = await graph.run(state=SimpleState())
     print(sorted(result))
     #> [1, 4, 9, 16, 25]
@@ -139,7 +140,7 @@ async def main():
     graph = g.build()
     result = await graph.run(state=SimpleState())
     print(result)
-    #> {'apple': 5, 'banana': 6, 'cherry': 6}
+    #> {'cherry': 6, 'banana': 6, 'apple': 5}
 ```
 
 _(This example is complete, it can be run "as is" — you'll need to add `import asyncio; asyncio.run(main())` to run `main`)_
@@ -410,7 +411,7 @@ _(This example is complete, it can be run "as is" — you'll need to add `import
 Like steps, joins can have custom IDs:
 
 ```python {title="join_custom_id.py" requires="basic_join.py"}
-from basic_join import g, ListReducer
+from basic_join import ListReducer, g
 
 my_join = g.join(ListReducer[int], node_id='my_custom_join_id')
 ```

docs/graph/beta/steps.md

@@ -16,21 +16,21 @@ from pydantic_graph.beta import GraphBuilder, StepContext
 class MyState:
     counter: int = 0
 
+g = GraphBuilder(state_type=MyState, output_type=int)
 
-async def main():
-    g = GraphBuilder(state_type=MyState, output_type=int)
+@g.step
+async def increment(ctx: StepContext[MyState, None, None]) -> int:
+    ctx.state.counter += 1
+    return ctx.state.counter
 
-    @g.step
-    async def increment(ctx: StepContext[MyState, None, None]) -> int:
-        ctx.state.counter += 1
-        return ctx.state.counter
+g.add(
+    g.edge_from(g.start_node).to(increment),
+    g.edge_from(increment).to(g.end_node),
+)
 
-    g.add(
-        g.edge_from(g.start_node).to(increment),
-        g.edge_from(increment).to(g.end_node),
-    )
+graph = g.build()
 
-    graph = g.build()
+async def main():
     state = MyState()
     result = await graph.run(state=state)
     print(result)
@@ -195,8 +195,11 @@ _(This example is complete, it can be run "as is" — you'll need to add `import
 By default, step node IDs are inferred from the function name. You can override this:
 
 ```python {title="custom_id.py" requires="basic_step.py"}
+from pydantic_graph.beta import StepContext
+
 from basic_step import MyState, g
 
+
 @g.step(node_id='my_custom_id')
 async def my_step(ctx: StepContext[MyState, None, None]) -> int:
     return 42
@@ -209,8 +212,11 @@ async def my_step(ctx: StepContext[MyState, None, None]) -> int:
 Labels provide documentation for diagram generation:
 
 ```python {title="labels.py" requires="basic_step.py"}
+from pydantic_graph.beta import StepContext
+
 from basic_step import MyState, g
 
+
 @g.step(label='Increment the counter')
 async def increment(ctx: StepContext[MyState, None, None]) -> int:
     ctx.state.counter += 1
@@ -330,6 +336,17 @@ The beta graph API provides strong type checking through generics. Type paramete
 - Input/output types match across edges
 
 ```python
+from dataclasses import dataclass
+
+from pydantic_graph.beta import GraphBuilder, StepContext
+
+
+@dataclass
+class MyState:
+    pass
+
+g = GraphBuilder(state_type=MyState, output_type=str)
+
 # Type checker will catch mismatches
 @g.step
 async def expects_int(ctx: StepContext[MyState, None, int]) -> str: