Skip to content

[Python] Do not fail type checks on Protocol type hints - #40757

Open
raams95 wants to merge 1 commit into
apache:masterfrom
raams95:fix-protocol-type-hint
Open

raams95 wants to merge 1 commit into
apache:masterfrom
raams95:fix-protocol-type-hint

Conversation

@raams95

@raams95 raams95 commented Oct 23, 2026

Copy link
Copy Markdown
Contributor

Fixes #20986

A type hint with a typing.Protocol that is not @runtime_checkable makes the pipeline fail at construction:

class HasUpper(typing.Protocol):
  def upper(self) -> str:
    ...

def fn(element: HasUpper) -> str:
  return element.upper()

with beam.Pipeline() as p:
  _ = p | beam.Create(['a', 'b']) | beam.Map(fn)
TypeError: Instance and class checks can only be used with @runtime_checkable protocols

is_consistent_with(str, HasUpper) calls issubclass, through beartype's is_subhint or directly, and Python raises this error for such a Protocol. With --runtime_type_check or --performance_runtime_type_check, check_constraint calls isinstance on each element and fails the same way. So a non-runtime Protocol hint works today only when the input type is exactly the same Protocol.

A Protocol is matched structurally, and Beam cannot check that with issubclass or isinstance. With this change, is_consistent_with returns True when the base is such a Protocol, and check_constraint skips it, as it does for Any. For these Protocols, both calls always raised the TypeError before. The only caller that caught it is the runtime check of a Union, which counted it as a mismatch, so a runtime check of Optional[HasUpper] now also passes for a str. The element type stays the Protocol, so the coder does not change. @runtime_checkable Protocols keep the issubclass/isinstance checks that they have today.

Tests:

  • test_typed_map_protocol_input in typed_pipeline_test.py runs the pipeline above.
  • test_protocol_type_hint in typecheck_test.py runs it with runtime_type_check=True.

Both fail before the change and pass after it on Python 3.11, 3.12, 3.13 and 3.14.

Testing:

  • pytest apache_beam/typehints/typed_pipeline_test.py apache_beam/typehints/typecheck_test.py -k protocol on Python 3.11, 3.12, 3.13 and 3.14
  • pytest apache_beam/typehints apache_beam/transforms apache_beam/coders apache_beam/pipeline_test.py on Python 3.12
  • yapf --diff and ruff check on the changed files

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 or fixes #456).
  • Update CHANGES.md with noteworthy changes.
  • If this contribution is large, please file an Apache Individual Contributor License Agreement. (Not needed, small change.)

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

A typing.Protocol that is not runtime_checkable makes issubclass and
isinstance raise TypeError, so a hint such as def fn(x: MyProtocol)
in beam.Map(fn) failed pipeline type checking unless the input type was
the same Protocol, and failed every element with runtime type checking.

Such a Protocol is matched structurally, which these checks cannot do,
so is_consistent_with now treats it as consistent and check_constraint
skips it.

Fixes apache#20986
@github-actions github-actions Bot added python core size/s S: 10-49 lines changed labels Oct 23, 2026

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

📋 Review Summary

This PR prevents pipeline construction and runtime type checks from raising TypeError when encountering non-@runtime_checkable typing.Protocol type hints by treating them as consistent in is_consistent_with and skipping isinstance in check_constraint. The approach is clean and well-targeted, with one ordering issue in is_consistent_with when the input type is a TypeConstraint.

🔍 General Feedback

  • _is_unchecked_protocol and its usage in check_constraint cleanly handle simple and nested composite hints at runtime.
  • In is_consistent_with, placing the _is_unchecked_protocol(base) check before the isinstance(sub, RowTypeConstraint) and isinstance(sub, TypeConstraint) branches will ensure protocols also work when the input PCollection has a Beam composite/container type (such as List[...], Tuple[...], Dict[...], or Row).

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

# Cannot check unsupported parameterized generic which will cause issubclass
# to fail with an exception.
return False
elif _is_unchecked_protocol(base):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 Medium: Because elif _is_unchecked_protocol(base): is placed after elif isinstance(sub, RowTypeConstraint): (line 1540) and elif isinstance(sub, TypeConstraint): (line 1542), is_consistent_with(sub, base) still returns False whenever sub is a Beam TypeConstraint (for example, List[int], Tuple[int, int], Dict[str, int], Set[int], or RowTypeConstraint—such as in p | beam.Create([[1, 2]]) | beam.Map(fn) where fn is annotated with a Protocol).

Moving the _is_unchecked_protocol(base) check above elif isinstance(sub, RowTypeConstraint): (or right after if isinstance(sub, AnyTypeConstraint) or isinstance(base, AnyTypeConstraint):) ensures that container/row PCollection element types are also considered consistent with an unchecked Protocol.

@codecov

codecov Bot commented Oct 23, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.89%. Comparing base (a8d6afd) to head (20b3f9e).
⚠️ Report is 70 commits behind head on master.

Additional details and impacted files
@@            Coverage Diff             @@
##           master   #40757      +/-   ##
==========================================
- Coverage   80.04%   79.89%   -0.16%     
==========================================
  Files         490      493       +3     
  Lines       82916    83469     +553     
==========================================
+ Hits        66373    66688     +315     
- Misses      16543    16781     +238     
Flag Coverage Δ
python 79.89% <100.00%> (-0.16%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @damccorm for label python.
Backup reviewer: tvalentyn (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

This pull request likely touches a core component ("core" label). Please review with scrutiny.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

  • stop reviewer notifications - opt out of the automated review tooling
  • remind me after tests pass - tag the comment author after tests pass
  • waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

beam crash when annotating function with typing.Protocol

1 participant