Skip to content

fix(pydantic): allow union of BaseModel types as stream_type#754

Open
Ghraven wants to merge 2 commits into
apache:mainfrom
Ghraven:fix/stream-type-union-basemodel
Open

fix(pydantic): allow union of BaseModel types as stream_type#754
Ghraven wants to merge 2 commits into
apache:mainfrom
Ghraven:fix/stream-type-union-basemodel

Conversation

@Ghraven
Copy link
Copy Markdown

@Ghraven Ghraven commented May 3, 2026

Fixes #607

What

Broadens the stream_type parameter type annotation in streaming_action.pydantic() and pydantic_streaming_action() to accept a union of pydantic BaseModel types in addition to a single Type[BaseModel] or Type[dict].

Why

Users who pass a union of models (e.g. MyModel1 | MyModel2) hit a static type error at the call site because the parameter was typed as Union[Type[BaseModel], Type[dict]], which only accepts a single concrete model type. The runtime validation logic in _validate_and_extract_signature_types_streaming() already handles union types correctly — this was a type annotation gap only.

Changes

  • burr/core/action.py: stream_type parameter in streaming_action.pydantic() widened to accept object (covers union types, which are not expressible as Type[X]).
  • burr/integrations/pydantic.py: Same change to PartialType alias and the validator function signature.
  • Docstring updated to mention union support.

Testing

No behaviour change — this is a type annotation fix only. Existing tests continue to pass. Users can now write:

union = MyModel1 | MyModel2

@streaming_action.pydantic(
    reads=[...],
    writes=[...],
    state_input_type=...,
    state_output_type=...,
    stream_type=union,
)

without a type error.

Ghraven added 2 commits May 3, 2026 16:56
Fixes apache#607

Broadens the `stream_type` parameter in `streaming_action.pydantic()` and
`pydantic_streaming_action()` to accept a union of pydantic BaseModel types
(e.g. `ModelA | ModelB`) in addition to a single `Type[BaseModel]` or
`Type[dict]`. The internal validation logic already handles this at runtime;
this change removes the overly strict static type annotation that rejected
union types at the call site.
Fixes apache#607

Broadens the `stream_type` parameter in `streaming_action.pydantic()` and
`pydantic_streaming_action()` to accept a union of pydantic BaseModel types
(e.g. `ModelA | ModelB`) in addition to a single `Type[BaseModel]` or
`Type[dict]`. The internal validation logic already handles this at runtime;
this change removes the overly strict static type annotation that rejected
union types at the call site.
@github-actions github-actions Bot added area/core Application, State, Graph, Actions area/integrations External integrations (LLMs, frameworks) area/typing Mypy, type hints, pydantic labels May 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/core Application, State, Graph, Actions area/integrations External integrations (LLMs, frameworks) area/typing Mypy, type hints, pydantic

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Streaming Event type, type hint, should support union type

1 participant