Add support for tagged output type hints. #37434

claudevdm · 2026-01-27T16:49:40Z

Please add a meaningful description for your change here

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

gemini-code-assist · 2026-01-27T16:53:52Z

Warning

Gemini is experiencing higher than usual traffic and was unable to create the summary. Please try again in a few hours by commenting /gemini summary.

claudevdm · 2026-01-28T17:25:10Z

/gemini summary

gemini-code-assist · 2026-01-28T17:25:35Z

Summary of Changes

This pull request significantly advances the type hinting capabilities within the Apache Beam Python SDK by introducing robust support for tagged outputs. It enables developers to explicitly define types for all outputs of a transform, whether main or side outputs, through decorator, method chaining, or function annotation styles. This enhancement aims to improve code readability, facilitate static analysis, and reduce runtime type-related errors in complex Beam pipelines that utilize multiple outputs.

Highlights

Enhanced Tagged Output Type Hinting: Introduced comprehensive support for type hints on tagged outputs across various Beam transforms (ParDo, Map, FlatMap, etc.). This allows users to specify types for both main and side outputs, improving type safety and static analysis.
Updated with_output_types API: The with_output_types method and decorator now accept keyword arguments to define type hints for specific tagged outputs (e.g., errors=str), making the API more expressive for multi-output scenarios.
Automatic Type Hint Extraction from Annotations: New internal helper functions were added to automatically parse and extract main and tagged output types from function annotations, including complex Union and Iterable types, simplifying type declaration for users.
Improved TaggedOutput Class: The TaggedOutput class has been made generic, allowing for more precise type annotations using TypeVar and Generic, which enhances the clarity and correctness of type definitions.
Warning for Undeclared Tags: The system now logs a warning if with_outputs requests tags that have not been explicitly declared in the type hints, helping to catch potential misconfigurations early.
New Dedicated Test Suite: A new test file (tagged_output_typehints_test.py) has been added to thoroughly validate the new functionality across all supported type hinting styles (decorator, method chain, and function annotation) and transforms.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

sdks/python/apache_beam/pvalue.py
- Added _tagged_output_types attribute to DoOutputsTuple to store type hints for tagged outputs.
- Updated PCollection.__getitem__ to use specific element_type from _tagged_output_types for tagged outputs, falling back to Any.
- Made TaggedOutput generic (TaggedOutput[TagType, ValueType]) for more precise type annotations.
sdks/python/apache_beam/transforms/core.py
- Implemented a warning in with_outputs if requested tags are not declared in type hints.
- Modified Map, MapTuple, and FlatMapTuple to pass tagged_output_types to with_output_types.
sdks/python/apache_beam/transforms/ptransform.py
- Updated with_output_types method to accept **tagged_type_hints for specifying types of tagged outputs.
- Added documentation and an example for using with_output_types with tagged outputs.
- Modified type_check_inputs_or_outputs to allow keyword arguments for output type hints, specifically for tagged outputs.
sdks/python/apache_beam/typehints/decorators.py
- Imported collections.abc, Literal, get_args, get_origin, and TaggedOutput.
- Introduced new helper functions: _tag_and_type, _contains_tagged_output, _extract_main_and_tagged, and _extract_output_types for parsing return annotations.
- Updated IOTypeHints.from_callable to use _extract_output_types and store tagged output types.
- Added IOTypeHints.tagged_output_types method and adjusted simple_output_type, has_simple_output_type, and strip_iterable to account for tagged outputs.
- Modified with_output_types decorator to accept **tagged_type_hints and updated its documentation.
sdks/python/apache_beam/typehints/decorators_test.py
- Added TaggedOutput import.
- Added new test cases for IOTypeHints.from_callable to verify correct extraction of main and tagged output types.
- Added a new test class TaggedOutputExtractionTest to specifically test the new helper functions for tagged output type hint extraction.
sdks/python/apache_beam/typehints/tagged_output_typehints_test.py
- Added new file containing extensive tests for tagged output type hinting across various transforms and type hinting styles (decorator, method chain, and function annotation).
sdks/python/apache_beam/typehints/typehints_test.py
- Updated the regex in test_no_kwargs_accepted to reflect the new behavior where kwargs are now accepted for tagged output types.

Activity

A bot attempted to summarize the PR but encountered high traffic and failed.
The author, claudevdm, manually requested a summary from the bot.

claudevdm · 2026-01-28T17:26:33Z

/gemini review

gemini-code-assist

Code Review

The pull request introduces comprehensive support for tagged output type hints in Apache Beam's Python SDK. This is a significant enhancement that improves type safety and clarity for multi-output transforms. The changes span across pvalue.py, transforms/core.py, and typehints/decorators.py, with corresponding updates to existing tests and the addition of a new, thorough test file (tagged_output_typehints_test.py). The implementation correctly handles the extraction, propagation, and validation of tagged output types through various mechanisms, including decorators, method chaining, and function annotations. The addition of warning messages for undeclared tags in with_outputs() is a good user experience improvement. The overall approach is well-structured and robust.

claudevdm · 2026-01-28T17:31:23Z

R: @damccorm

claudevdm · 2026-01-28T17:31:30Z

R: @robertwb

claudevdm · 2026-01-28T17:31:36Z

R: @jrmccluskey

github-actions · 2026-01-28T17:32:34Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

github-actions · 2026-01-28T17:32:43Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

github-actions · 2026-01-28T17:32:49Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

jrmccluskey

LGTM, thanks! I like where this wound up, should be pretty easy for users to add to their pipelines

sdks/python/apache_beam/typehints/decorators.py

claudevdm · 2026-01-28T19:51:23Z

sdks/python/apache_beam/typehints/decorators.py

+    return None, {tag: typ}
+
+  if t is TaggedOutput:
+    raise TypeError(


Actually, should this fail or just warn? I saw an issue #26493 where someone is using typehint -> Iterable[Union[Dict[str, Any], pvalue.TaggedOutput]]:

Their pipeline will break until they fix the typehint. Maybe we should rather warn that the tagged typehints cannot be parsed instead of throwing

We should probably allow this as pvalue.TaggedOutput[*, Any], though perhaps with a warning.

I'm thinking for the sake of update compat, it makes sense to fall back to old behavior for bare TaggedOutput. This means the main and tagged outputs will fall back to Union -> Any -> FastPrimitivesCoder?

If we allow this as pvalue.TaggedOutput[*, Any], there is a chance the main output will now map to a different coder e.g.

def dofn(element) -> int | TaggedOutput

If we still extract main, and map tagged outputs to Any then main type will change from Union[int, TaggedOutput] to int and the coder will change from FastPrimitivesCoder to IntCoder

That is generally a good thing. But if needed we could guard this with the update compatibility flag.

Do you think an update compat flag is necessary? If the coder changed from FastPrimitivesCoder to IntCoder, then updating a pipeline will fail. The user will just have to change their typehints to -> Any and then the update should work.

robertwb

Very nice.

Have you looked into updating type inference to be able to handle TaggedOutput types?

sdks/python/apache_beam/transforms/core.py

sdks/python/apache_beam/transforms/ptransform.py

robertwb · 2026-01-29T20:25:08Z

sdks/python/apache_beam/typehints/decorators.py

+    return None, {tag: typ}
+
+  if t is TaggedOutput:
+    raise TypeError(


We should probably allow this as pvalue.TaggedOutput[*, Any], though perhaps with a warning.

sdks/python/apache_beam/typehints/decorators.py

robertwb · 2026-01-29T20:31:37Z

sdks/python/apache_beam/typehints/decorators.py

+    yield_type = get_args(return_annotation)[0]
+    clean_yield, tagged_types = _extract_main_and_tagged(yield_type)
+    clean_main = clean_yield if clean_yield else Any
+    return [Iterable[clean_main]], tagged_types


Seems a bit asymmetric here to be wrapping the main type in an iterable but not the tagged types. Should we instead be calling this after the (higher level) iterable unwrapping?

Agreed it is not ideal. I'm not sure what is better. Are you suggesting we defer tagged output extraction until during/after strip_iterable?

from_callable doesnt extract any tagged inputs, so

fn(x: int) -> Iterable[int | TaggedOutput[Literal['errors'], str]] -> output_types = ((Iterable[int | TaggedOutput[Literal['errors'], str]],), {})

convert_to_beam_type passes through TaggedOutput

convert_to_beam_type(Iterable[int | TaggedOutput[Literal['errors'], str]]) -> Iterable[Union[int, TaggedOutput[Literal['errors'], str]]] (TaggedOutput preserved)

strip_iterable() unwraps

Iterable[Union[int, TaggedOutput[...]]] → Union[int, TaggedOutput[...]]

Finally we extract tagged outputs?

robertwb · 2026-01-29T21:37:29Z

sdks/python/apache_beam/typehints/tagged_output_typehints_test.py

+
+class AnnotationStyleTaggedOutputTest(unittest.TestCase):
+  """Tests for function annotation style across all transforms."""
+  def test_map_annotation_union(self):


claudevdm · 2026-01-29T22:04:03Z

Very nice.

Have you looked into updating type inference to be able to handle TaggedOutput types?

What do you mean by this? Are you referring to

beam/sdks/python/apache_beam/pipeline.py

Line 995 in 96edda4

elif isinstance(result_pcollection, pvalue.DoOutputsTuple):

?

This already works and I should remove the TODO comment. The element type will already be set in DoOutputsTuple when the tagged pcollection is created.

github-actions bot added the python label Jan 27, 2026

Add tagged typehint support.

a5fcdc0

claudevdm force-pushed the kwarg-output-hints branch from 1bdd579 to a5fcdc0 Compare January 28, 2026 16:00

gemini-code-assist bot reviewed Jan 28, 2026

View reviewed changes

claudevdm requested a review from robertwb January 28, 2026 17:30

claudevdm marked this pull request as ready for review January 28, 2026 17:30

claudevdm requested review from damccorm and jrmccluskey January 28, 2026 17:31

jrmccluskey approved these changes Jan 28, 2026

View reviewed changes

claudevdm commented Jan 28, 2026

View reviewed changes

sdks/python/apache_beam/typehints/decorators.py Show resolved Hide resolved

claudevdm commented Jan 28, 2026

View reviewed changes

robertwb reviewed Jan 29, 2026

View reviewed changes

claudevdm added 3 commits January 29, 2026 17:20

Just warn when bare tagged output

4039780

Remove contains tagged output check.

5692c12

Mapped bare TaggedOutput to Any

c4e4d10

Add support for tagged output type hints. #37434

Are you sure you want to change the base?

Add support for tagged output type hints. #37434

Conversation

claudevdm commented Jan 27, 2026

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist bot commented Jan 27, 2026

Uh oh!

claudevdm commented Jan 28, 2026

Uh oh!

gemini-code-assist bot commented Jan 28, 2026

Summary of Changes

Highlights

Uh oh!

claudevdm commented Jan 28, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

claudevdm commented Jan 28, 2026

Uh oh!

claudevdm commented Jan 28, 2026

Uh oh!

claudevdm commented Jan 28, 2026

Uh oh!

github-actions bot commented Jan 28, 2026

Uh oh!

github-actions bot commented Jan 28, 2026

Uh oh!

github-actions bot commented Jan 28, 2026

Uh oh!

jrmccluskey left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

robertwb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claudevdm commented Jan 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants