Implements Support to Control Direction of State Diagram Generated by Mermaid Code (#716)

Co-authored-by: Sydney Runkle <[email protected]>

Author: izzyacademy

docs/graph.md

@@ -790,3 +790,35 @@ stateDiagram-v2
 classDef highlighted fill:#fdff32
 class Answer highlighted
 ```
+
+### Setting Direction of the State Diagram
+
+You can specify the direction of the state diagram using one of the following values:
+
+- `'TB'`: Top to bottom, the diagram flows vertically from top to bottom.
+- `'LR'`: Left to right, the diagram flows horizontally from left to right.
+- `'RL'`: Right to left, the diagram flows horizontally from right to left.
+- `'BT'`: Bottom to top, the diagram flows vertically from bottom to top.
+
+Here is an example of how to do this using 'Left to Right' (LR) instead of the default 'Top to Bottom' (TB)
+```py {title="vending_machine_diagram.py" py="3.10"}
+from vending_machine import InsertCoin, vending_machine_graph
+
+vending_machine_graph.mermaid_code(start_node=InsertCoin, direction='LR')
+```
+
+```mermaid
+---
+title: vending_machine_graph
+---
+stateDiagram-v2
+  direction LR
+  [*] --> InsertCoin
+  InsertCoin --> CoinsInserted
+  CoinsInserted --> SelectProduct
+  CoinsInserted --> Purchase
+  SelectProduct --> Purchase
+  Purchase --> InsertCoin
+  Purchase --> SelectProduct
+  Purchase --> [*]
+```

pydantic_graph/pydantic_graph/graph.py

@@ -286,6 +286,7 @@ def mermaid_code(
         highlighted_nodes: Sequence[mermaid.NodeIdent] | mermaid.NodeIdent | None = None,
         highlight_css: str = mermaid.DEFAULT_HIGHLIGHT_CSS,
         infer_name: bool = True,
+        direction: mermaid.StateDiagramDirection | None = None,
     ) -> str:
         """Generate a diagram representing the graph as [mermaid](https://mermaid.js.org/) diagram.
 
@@ -299,6 +300,7 @@ def mermaid_code(
             highlighted_nodes: Optional node or nodes to highlight.
             highlight_css: The CSS to use for highlighting nodes.
             infer_name: Whether to infer the graph name from the calling frame.
+            direction: The direction of flow.
 
         Returns:
             The mermaid code for the graph, which can then be rendered as a diagram.
@@ -346,6 +348,7 @@ def mermaid_code(
             title=title or None,
             edge_labels=edge_labels,
             notes=notes,
+            direction=direction,
         )
 
     def mermaid_image(

pydantic_graph/pydantic_graph/mermaid.py

@@ -16,11 +16,27 @@
 if TYPE_CHECKING:
     from .graph import Graph
 
-
-__all__ = 'NodeIdent', 'DEFAULT_HIGHLIGHT_CSS', 'generate_code', 'MermaidConfig', 'request_image', 'save_image'
+__all__ = (
+    'NodeIdent',
+    'DEFAULT_HIGHLIGHT_CSS',
+    'generate_code',
+    'MermaidConfig',
+    'request_image',
+    'save_image',
+    'StateDiagramDirection',
+)
 DEFAULT_HIGHLIGHT_CSS = 'fill:#fdff32'
 """The default CSS to use for highlighting nodes."""
 
+StateDiagramDirection = Literal['TB', 'LR', 'RL', 'BT']
+"""Used to specify the direction of the state diagram generated by mermaid.
+
+- `'TB'`: Top to bottom, this is the default for mermaid charts.
+- `'LR'`: Left to right
+- `'RL'`: Right to left
+- `'BT'`: Bottom to top
+"""
+
 
 def generate_code(  # noqa: C901
     graph: Graph[Any, Any, Any],
@@ -32,6 +48,7 @@ def generate_code(  # noqa: C901
     title: str | None = None,
     edge_labels: bool = True,
     notes: bool = True,
+    direction: StateDiagramDirection | None,
 ) -> str:
     """Generate [Mermaid state diagram](https://mermaid.js.org/syntax/stateDiagram.html) code for a graph.
 
@@ -43,6 +60,8 @@ def generate_code(  # noqa: C901
         title: The title of the diagram.
         edge_labels: Whether to include edge labels in the diagram.
         notes: Whether to include notes in the diagram.
+        direction: The direction of flow.
+
 
     Returns:
         The Mermaid code for the graph.
@@ -56,6 +75,8 @@ def generate_code(  # noqa: C901
     if title:
         lines = ['---', f'title: {title}', '---']
     lines.append('stateDiagram-v2')
+    if direction is not None:
+        lines.append(f'  direction {direction}')
     for node_id, node_def in graph.node_defs.items():
         # we use round brackets (rounded box) for nodes other than the start and end
         if node_id in start_node_ids:
@@ -131,6 +152,7 @@ def request_image(
         title=kwargs.get('title'),
         edge_labels=kwargs.get('edge_labels', True),
         notes=kwargs.get('notes', True),
+        direction=kwargs.get('direction'),
     )
     code_base64 = base64.b64encode(code.encode()).decode()
 
@@ -245,6 +267,8 @@ class MermaidConfig(TypedDict, total=False):
     """
     httpx_client: httpx.Client
     """An HTTPX client to use for requests, mostly for testing purposes."""
+    direction: StateDiagramDirection
+    """The direction of the state diagram."""
 
 
 NodeIdent: TypeAlias = 'type[BaseNode[Any, Any, Any]] | BaseNode[Any, Any, Any] | str'

tests/graph/test_mermaid.py

@@ -194,6 +194,38 @@ def test_mermaid_code_all_nodes():
 """)
 
 
+def test_mermaid_code_all_nodes_no_direction():
+    assert graph3.mermaid_code() == snapshot("""\
+---
+title: graph3
+---
+stateDiagram-v2
+  AllNodes --> AllNodes
+  AllNodes --> Foo
+  AllNodes --> Bar
+  Foo --> Bar
+  Bar --> [*]\
+""")
+
+
+def test_mermaid_code_all_nodes_with_direction_lr():
+    assert graph3.mermaid_code(direction='LR') == snapshot("""\
+---
+title: graph3
+---
+stateDiagram-v2
+  direction LR
+  AllNodes --> AllNodes
+  AllNodes --> Foo
+  AllNodes --> Bar
+  Foo --> Bar
+  Bar --> [*]\
+""")
+
+
+# Tests for direction ends here
+
+
 def test_docstring_notes_classvar():
     assert Spam.docstring_notes is True
     assert repr(Spam()) == 'Spam()'