diff --git a/burr/core/action.py b/burr/core/action.py index eb6fceff..92144b07 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -24,6 +24,16 @@ import textwrap import types import typing + +# types.UnionType was added in Python 3.10 to represent PEP 604 `X | Y` +# syntax. Burr supports Python >= 3.9, so on 3.9 we fall back to a sentinel +# type that keeps Union annotations well-formed. No PEP 604 union can ever +# exist on 3.9, so isinstance() checks against it simply never match. +if sys.version_info >= (3, 10): + _UnionType = types.UnionType +else: # pragma: no cover - exercised on Python 3.9 CI only + class _UnionType: # type: ignore[no-redef] + """Placeholder for ``types.UnionType`` on Python < 3.10.""" from collections.abc import AsyncIterator from typing import ( TYPE_CHECKING, @@ -1511,7 +1521,7 @@ def pydantic( writes: List[str], state_input_type: Type["BaseModel"], state_output_type: Type["BaseModel"], - stream_type: Union[Type["BaseModel"], Type[dict]], + stream_type: Union[Type["BaseModel"], Type[dict], _UnionType], tags: Optional[List[str]] = None, ) -> Callable: """Creates a streaming action that uses pydantic models. diff --git a/burr/integrations/pydantic.py b/burr/integrations/pydantic.py index 300cbb6a..3ef2a2ce 100644 --- a/burr/integrations/pydantic.py +++ b/burr/integrations/pydantic.py @@ -41,6 +41,7 @@ from burr.core.action import ( FunctionBasedAction, FunctionBasedStreamingAction, + _UnionType, bind, derive_inputs_from_fn, ) @@ -269,7 +270,7 @@ async def async_action_function(state: State, **kwargs) -> State: return decorator -PartialType = Union[Type[pydantic.BaseModel], Type[dict]] +PartialType = Union[Type[pydantic.BaseModel], Type[dict], _UnionType] PydanticStreamingActionFunctionSync = Callable[ ..., Generator[Tuple[Union[pydantic.BaseModel, dict], Optional[pydantic.BaseModel]], None, None] @@ -290,7 +291,7 @@ async def async_action_function(state: State, **kwargs) -> State: def _validate_and_extract_signature_types_streaming( fn: PydanticStreamingActionFunction, - stream_type: Optional[Union[Type[pydantic.BaseModel], Type[dict]]], + stream_type: Optional[Union[Type[pydantic.BaseModel], Type[dict], _UnionType]], state_input_type: Optional[Type[pydantic.BaseModel]] = None, state_output_type: Optional[Type[pydantic.BaseModel]] = None, ) -> Tuple[