fix tests for other versions

pydantic · samuelcolvin · Jan 15, 2025 · Dec 22, 2024 · Dec 22, 2024 · Dec 22, 2024
commit 61f3d2b7710ee0e943526be85b9e8e11a1101e0d
diff --git a/pydantic_ai_graph/pydantic_ai_graph/_utils.py b/pydantic_ai_graph/pydantic_ai_graph/_utils.py
@@ -1,6 +1,10 @@
+from __future__ import annotations as _annotations
+
 import sys
 import types
-from typing import Any, TypeAliasType, Union, get_args, get_origin
+from typing import Any, Union, get_args, get_origin
+
+from typing_extensions import TypeAliasType
 
 
 def get_union_args(tp: Any) -> tuple[Any, ...]:

diff --git a/tests/test_graph.py b/tests/test_graph.py
@@ -1,6 +1,7 @@
 from __future__ import annotations as _annotations
 
 from datetime import timezone
+from typing import Union
 
 import pytest
 from inline_snapshot import snapshot
@@ -22,7 +23,7 @@ async def run(self, ctx: GraphContext) -> Double:
             return Double(len(self.input_data))
 
     class Double(BaseNode[None, int, int]):
-        async def run(self, ctx: GraphContext) -> String2Length | End[int]:
+        async def run(self, ctx: GraphContext) -> Union[String2Length, End[int]]:  # noqa: UP007
             if self.input_data == 7:
                 return String2Length('x' * 21)
             else: