Skip to content

fix(pydantic): accept union type objects for stream_type hints #617

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
majiayu000 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(pydantic): accept union type objects for stream_type hints #617

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion burr/core/action.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,18 @@
except ImportError:
pass

try:
UnionType = types.UnionType
except AttributeError: # py<3.10
UnionType = None

if typing.TYPE_CHECKING:
StreamingPydanticType = Any
else:
StreamingPydanticType = Union[Type["BaseModel"], Type[dict]]
if UnionType is not None:
StreamingPydanticType = Union[Type["BaseModel"], Type[dict], UnionType]
Comment on lines +70 to +75
Copy link
Copy Markdown
Contributor

@skrawcz skrawcz Jan 4, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

doing Any for typing checking will of course work and is cheating...

I think we shouldn't have a typing.TYPE_CHECKING part here, and instead I think the fix is line 74 + 75.

note I had to do this stream_type: Union[Type["BaseModel"], Type[dict], type, "types.UnionType"], for my pyright in the IDE to not complain -- i.e. the quotes part.

Copy link
Copy Markdown
Contributor

@skrawcz skrawcz Apr 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The TYPE_CHECKING = Any pattern effectively disables the type checker for stream_type. Replace with:

if sys.version_info >= (3, 10):
    StreamingPydanticType = Union[Type["BaseModel"], Type[dict], types.UnionType]
else:
    StreamingPydanticType = Union[Type["BaseModel"], Type[dict]]

sys is already imported in this file. Pyright natively narrows on sys.version_info checks.



class Function(abc.ABC):
"""Interface to represent the 'computing' part of an action"""
Expand Down Expand Up @@ -1309,7 +1321,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: "StreamingPydanticType",
tags: Optional[List[str]] = None,
) -> Callable:
"""Creates a streaming action that uses pydantic models.
Expand Down
19 changes: 14 additions & 5 deletions burr/integrations/pydantic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import types
import typing
from typing import (
Copy link
Copy Markdown
Contributor

@skrawcz skrawcz Apr 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: the Any import can be removed if switching to the sys.version_info approach.

Any,
AsyncGenerator,
Awaitable,
Callable,
Expand Down Expand Up @@ -264,7 +265,17 @@ async def async_action_function(state: State, **kwargs) -> State:
return decorator


PartialType = Union[Type[pydantic.BaseModel], Type[dict]]
try:
UnionType = types.UnionType
except AttributeError: # py<3.10
UnionType = None

if typing.TYPE_CHECKING:
PartialType = Any
else:
PartialType = Union[Type[pydantic.BaseModel], Type[dict]]
if UnionType is not None:
PartialType = Union[Type[pydantic.BaseModel], Type[dict], UnionType]
Copy link
Copy Markdown
Contributor

@skrawcz skrawcz Apr 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same issue — replace the TYPE_CHECKING = Any pattern with the sys.version_info guard for PartialType.


PydanticStreamingActionFunctionSync = Callable[
..., Generator[Tuple[Union[pydantic.BaseModel, dict], Optional[pydantic.BaseModel]], None, None]
Expand All @@ -285,12 +296,10 @@ 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[PartialType],
state_input_type: Optional[Type[pydantic.BaseModel]] = None,
state_output_type: Optional[Type[pydantic.BaseModel]] = None,
) -> Tuple[
Type[pydantic.BaseModel], Type[pydantic.BaseModel], Union[Type[dict], Type[pydantic.BaseModel]]
]:
) -> Tuple[Type[pydantic.BaseModel], Type[pydantic.BaseModel], PartialType]:
if stream_type is None:
# TODO -- derive from the signature
raise ValueError(f"stream_type is required for function: {fn.__qualname__}")
Expand Down
Loading