[Python] Do not fail type checks on Protocol type hints

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

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):

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)]

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)