Skip to content

fix: allow union types (X | Y) for stream_type in streaming_action.pydantic#845

Open
mirakour wants to merge 3 commits into
apache:mainfrom
mirakour:fix/stream-type-union-support
Open

fix: allow union types (X | Y) for stream_type in streaming_action.pydantic#845
mirakour wants to merge 3 commits into
apache:mainfrom
mirakour:fix/stream-type-union-support

Conversation

@mirakour

Copy link
Copy Markdown

Closes #607

What and Why

The stream_type parameter in @streaming_action.pydantic and pydantic_streaming_action currently only accepts a single Pydantic model class or dict. Passing a union type (e.g. ModelA | ModelB using Python 3.10+ syntax) raises a TypeError because types.UnionType is not included in the PartialType alias.

This PR fixes that by expanding PartialType to include types.UnionType on Python 3.10+ using a version guard, and updating the stream_type annotation in action.py accordingly.

Changes

  • burr/integrations/pydantic.py: Added import sys and a version guard so PartialType includes types.UnionType on Python 3.10+
  • burr/core/action.py: Updated the stream_type annotation in streaming_action.pydantic() from Union[Type["BaseModel"], Type[dict]] to Union[Type["BaseModel"], Type[dict], Any]
  • tests/integrations/test_burr_pydantic.py: Added 3 regression tests guarded with @pytest.mark.skipif(sys.version_info < (3, 10), ...)

How I tested this

Added 3 regression tests in tests/integrations/test_burr_pydantic.py:

  1. test_validate_streaming_signature_accepts_union_stream_type — verifies the internal validation function accepts types.UnionType
  2. test_pydantic_streaming_action_accepts_union_stream_type — verifies @pydantic_streaming_action works with stream_type=ModelA | ModelB
  3. test_streaming_action_pydantic_decorator_accepts_union_stream_type — same for the @streaming_action.pydantic decorator

All tests are skipped automatically on Python < 3.10.

Notes

  • No behavior change on Python < 3.10; the original PartialType is preserved for older versions
  • The Any annotation in action.py is intentionally broad to accommodate future union-like types without requiring another version guard

Checklist

  • PR has an informative and human-readable title (this will be pulled into the release notes)
  • Changes are limited to a single goal (no scope creep)
  • Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
  • Any change in functionality is tested
  • New functions are documented (with a description, list of inputs, and expected output)
  • Placeholder code is flagged / future TODOs are captured in comments
  • Project documentation has been updated if adding/changing functionality.

mirakour added 3 commits July 3, 2026 14:39
Add regression tests for union stream type support in Pydantic decorators. These tests ensure compatibility with Python 3.10's union type syntax.
@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 Jul 22, 2026
@eeshsaxena

Copy link
Copy Markdown

Heads up that #813 (opened 2026-06-20) implements the same feature, touching the same two files. Worth the maintainers deciding between them rather than reviewing both in isolation.

Having read both, there are two cases this one may not cover. requires-python is >=3.9, so both matter:

1. typing.Union[A, B] on Python 3.10-3.13. The check here keys on types.UnionType, which only covers the PEP 604 A | B form. On 3.10-3.13 typing.Union[A, B] is a typing._UnionGenericAlias and is not an instance of types.UnionType, so it would still be rejected. Python 3.14 unified the two (I confirmed isinstance(Union[int, str], types.UnionType) is True there, and type(int | str) is type(Union[int, str])), but on the older supported versions they are distinct. The issue example uses |, though Union[...] seems worth accepting too since the existing annotation is written that way.

2. Union members are not validated. Widening PartialType means int | str is accepted as a stream_type even though neither member is a BaseModel or dict.

#813 handles both by keeping the annotation wide and doing the real check at decoration time:

if typing.get_origin(t) is Union:                    # typing.Union[...]
    return all(_is_valid_stream_type(arg) for arg in typing.get_args(t))
_UnionType = getattr(types, "UnionType", None)        # X | Y, guarded for 3.9
if _UnionType is not None and isinstance(t, _UnionType):
    return all(_is_valid_stream_type(arg) for arg in t.__args__)

That also speaks to @elijahbenizzy's question on #607 ("does it actually detect per-instance rather than assume they're all the same?"): recursing into the members is what makes an invalid union fail at decoration time rather than surfacing later during streaming.

Not trying to talk down your PR - the smaller diff is appealing and the version guard is careful. But if the two get merged into one, the member validation and the typing.Union path look worth carrying over.

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

2 participants