Skip to content

Commit 812a469

Browse files
authored
Make Edge hashable (#2064)
1 parent 0d11906 commit 812a469

File tree

2 files changed

+15
-2
lines changed

2 files changed

+15
-2
lines changed

pydantic_graph/pydantic_graph/nodes.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ def generate_snapshot_id(node_id: str) -> str:
174174
return f'{node_id}:{uuid4().hex}'
175175

176176

177-
@dataclass
177+
@dataclass(frozen=True)
178178
class Edge:
179179
"""Annotation to apply a label to an edge in a graph."""
180180

tests/graph/test_mermaid.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from dataclasses import dataclass
66
from datetime import timezone
77
from pathlib import Path
8-
from typing import Annotated, Callable
8+
from typing import Annotated, Callable, Union
99

1010
import httpx
1111
import pytest
@@ -426,3 +426,16 @@ async def run(self, ctx: GraphRunContext) -> int: # type: ignore
426426

427427
with pytest.raises(GraphSetupError, match="Invalid return type: <class 'int'>"):
428428
NoReturnType.get_node_def({})
429+
430+
431+
def test_edge_union():
432+
"""Test that a union of things annotated with an Edge doesn't raise a TypeError.
433+
434+
This is important because such unions may occur as a return type for a graph, and needs to be evaluated when
435+
generating a mermaid diagram.
436+
"""
437+
# This would raise an error on 3.10 if Edge was not hashable:
438+
edges_union = Union[
439+
Annotated[End[None], Edge(label='first label')], Annotated[End[None], Edge(label='second label')]
440+
]
441+
assert edges_union

0 commit comments

Comments
 (0)