From 20b3f9eda50d585452704bb4e8ac28338aeb2918 Mon Sep 17 00:00:00 2001 From: JanakiRamChittineni <341781521+raams95@users.noreply.github.com> Date: Fri, 23 Oct 2026 16:11:09 -0500 Subject: [PATCH] [Python] Do not fail type checks on Protocol type hints 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 #20986 --- CHANGES.md | 1 + .../apache_beam/typehints/typecheck_test.py | 14 ++++++++++++++ .../apache_beam/typehints/typed_pipeline_test.py | 13 +++++++++++++ sdks/python/apache_beam/typehints/typehints.py | 15 +++++++++++++++ 4 files changed, 43 insertions(+) diff --git a/CHANGES.md b/CHANGES.md index c9350da11d69..7831ca046060 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -138,6 +138,7 @@ * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)). * (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)). * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)). +* (Python) Type hints with a `typing.Protocol` that is not `@runtime_checkable`, such as `def fn(x: MyProtocol)` in `beam.Map(fn)`, no longer make pipeline or runtime type checking fail with a `TypeError` ([#20986](https://github.com/apache/beam/issues/20986)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/python/apache_beam/typehints/typecheck_test.py b/sdks/python/apache_beam/typehints/typecheck_test.py index 81046cd3cc0f..c71448e0ceba 100644 --- a/sdks/python/apache_beam/typehints/typecheck_test.py +++ b/sdks/python/apache_beam/typehints/typecheck_test.py @@ -26,6 +26,7 @@ # ruff: noqa: UP006 import os import tempfile +import typing import unittest from typing import Iterable from typing import Tuple @@ -134,6 +135,19 @@ def test_wrapper_pipeline_type_check(self): with self.assertRaisesRegex(ValueError, r'int.*is not iterable'): _ = self.p | beam.Create([1, 2, 3]) | beam.ParDo(dofn) + def test_protocol_type_hint(self): + # https://github.com/apache/beam/issues/20986 + class HasUpper(typing.Protocol): + def upper(self) -> str: + ... + + def fn(e: HasUpper) -> str: + return e.upper() + + result = self.p | beam.Create(['a', 'b']) | beam.Map(fn) + assert_that(result, equal_to(['A', 'B'])) + self.p.run() + class PerformanceRuntimeTypeCheckTest(unittest.TestCase): def setUp(self): diff --git a/sdks/python/apache_beam/typehints/typed_pipeline_test.py b/sdks/python/apache_beam/typehints/typed_pipeline_test.py index 9a5efd3dead7..a65cc6a0309e 100644 --- a/sdks/python/apache_beam/typehints/typed_pipeline_test.py +++ b/sdks/python/apache_beam/typehints/typed_pipeline_test.py @@ -462,6 +462,19 @@ def fn(element: int) -> typehints.Optional[int]: result = [1, 2, 3] | beam.Map(fn) self.assertCountEqual([None, 2, 3], result) + def test_typed_map_protocol_input(self): + # https://github.com/apache/beam/issues/20986 + class HasUpper(typing.Protocol): + def upper(self) -> str: + ... + + def fn(element: HasUpper) -> str: + return element.upper() + + with TestPipeline() as p: + result = p | beam.Create(['a', 'b']) | beam.Map(fn) + assert_that(result, equal_to(['A', 'B'])) + def test_inferred_type_of_tuple_slice(self): # https://github.com/apache/beam/issues/20899 rows = [('day1', 'a', 1), ('day2', 'a', 2), ('day1', 'b', 3)] diff --git a/sdks/python/apache_beam/typehints/typehints.py b/sdks/python/apache_beam/typehints/typehints.py index 041a719ff831..c639ad498316 100644 --- a/sdks/python/apache_beam/typehints/typehints.py +++ b/sdks/python/apache_beam/typehints/typehints.py @@ -368,6 +368,16 @@ def is_typing_generic(type_param): return isinstance(type_param, typing._GenericAlias) +def _is_unchecked_protocol(type_param): + """Determines if an object is a typing.Protocol class that is not + runtime_checkable, which isinstance() and issubclass() reject. + """ + return ( + isinstance(type_param, type) and + getattr(type_param, '_is_protocol', False) and + not getattr(type_param, '_is_runtime_protocol', False)) + + def validate_composite_type_param(type_param, error_msg_prefix): """Determines if an object is a valid type parameter to a :class:`CompositeTypeHint`. @@ -429,6 +439,8 @@ def check_constraint(type_constraint, object_instance): pass elif not isinstance(type_constraint, type): raise RuntimeError("bad type: %s" % (type_constraint, )) + elif _is_unchecked_protocol(type_constraint): + pass elif not isinstance(object_instance, type_constraint): raise SimpleTypeHintError @@ -1536,6 +1548,9 @@ def is_consistent_with( # Cannot check unsupported parameterized generic which will cause issubclass # to fail with an exception. return False + elif _is_unchecked_protocol(base): + # Protocols are matched structurally, which issubclass cannot check. + return True if use_beartype: try: return is_subhint(sub, base)