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

[mirakour]

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.

[eeshsaxena]

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.

[eeshsaxena]

I tried these tests against an unmodified main and all three pass there, so I don't think they guard the fix.

_validate_and_extract_signature_types_streaming never inspects stream_type at runtime. It checks the three arguments for None and returns stream_type unchanged, so a union already goes through today:

@pydantic_streaming_action(
    reads=["count"], writes=["count"],
    state_input_type=St, state_output_type=St,
    stream_type=ModelA | ModelB,   # no error on main
)
def act(state: St): ...

stream_type=42 is accepted on main too. I appended the three new tests to tests/integrations/test_burr_pydantic.py at 805131a and got 3 passed. That fits #607, which is about the type hint being rejected by a type checker rather than anything failing at runtime, so the only thing that can catch a regression here is a typing check (mypy over a small example), not pytest.

Separately, the annotation in burr/core/action.py:

stream_type: Union[Type["BaseModel"], Type[dict], Any],

Any absorbs the other members of a Union, so this collapses to plain Any and the parameter ends up with no type at all, stream_type=42 included. The version-gated PartialType you added in integrations/pydantic.py is the precise version of this. Importing and using it in action.py would keep the hint meaningful.

One more case: types.UnionType only covers the A | B spelling. typing.Union[ModelA, ModelB] is still valid on 3.10+ and is what people writing 3.9-compatible code use, and it would still be flagged. Worth adding to the alias.