[Python] Convert typing and native generic hints in Watch coder inference (#39547)

registry.get_coder receives typing and native generic annotations such
as tuple[str, float] unconverted and falls back to pickling. Watch now
converts hints with convert_to_beam_type before the registry lookup, so
annotated poll functions and key functions infer real structured coders
and fully typed elements.

Author: Eliaaazzz

sdks/python/apache_beam/io/watch.py

@@ -77,6 +77,7 @@ def poll(prefix) -> PollResult[str]:
 from apache_beam.transforms import PTransform
 from apache_beam.transforms import core
 from apache_beam.transforms.window import TimestampedValue
+from apache_beam.typehints import native_type_compatibility
 from apache_beam.utils.timestamp import MAX_TIMESTAMP
 from apache_beam.utils.timestamp import Duration
 from apache_beam.utils.timestamp import Timestamp
@@ -642,6 +643,13 @@ def _poll_output_type(poll_fn) -> Any:
   return Any
 
 
+def _coder_for_hint(hint) -> Coder:
+  # typing and native generic hints such as tuple[str, float] must be
+  # converted to Beam typehints, or the registry falls back to pickling.
+  return coders.registry.get_coder(
+      native_type_compatibility.convert_to_beam_type(hint))
+
+
 class Watch(PTransform):
   """Watches a growing set of outputs per input via a periodic poll function.
 
@@ -689,14 +697,14 @@ def expand(self, pcoll):
     if output_coder is None and isinstance(self._poll_fn, PollFn):
       output_coder = self._poll_fn.default_output_coder()
     if output_coder is None:
-      output_coder = coders.registry.get_coder(_poll_output_type(self._poll_fn))
+      output_coder = _coder_for_hint(_poll_output_type(self._poll_fn))
     if self._output_key_fn is None:
       # The output is its own dedup key, so the key coder is the output coder.
       key_fn = _identity
       key_coder = self._output_key_coder or output_coder
     else:
       key_fn = self._output_key_fn
-      key_coder = self._output_key_coder or coders.registry.get_coder(
+      key_coder = self._output_key_coder or _coder_for_hint(
           _return_type(self._output_key_fn))
     # Dedup hashes the encoded key, so equal keys must encode equally; use the
     # coder's deterministic form and reject coders that have none.

sdks/python/apache_beam/io/watch_test.py

@@ -18,6 +18,7 @@
 """Tests for the Watch transform."""
 
 import collections
+import typing
 import unittest
 
 import apache_beam as beam
@@ -444,6 +445,24 @@ def test_infers_output_coder_from_return_annotation(self):
           | Watch(_complete_poll, poll_interval=Duration(1)))
       self.assertEqual(typehints.Tuple[str, str], output.element_type)
 
+  def test_infers_coder_from_generic_annotations(self):
+    # tuple[str, float] and typing.Tuple[str, float] resolve to a tuple coder,
+    # not the pickling fallback.
+    def native_poll(element) -> PollResult[tuple[str, float]]:
+      return PollResult.complete([(element, 1.0)])
+
+    def typing_poll(
+        element) -> PollResult[typing.Tuple[str, float]]:  # noqa: UP006
+      return PollResult.complete([(element, 1.0)])
+
+    for poll in (native_poll, typing_poll):
+      with self._in_memory_pipeline() as p:
+        output = (
+            p | beam.Create(['k:']) | Watch(poll, poll_interval=Duration(1)))
+        self.assertEqual(
+            typehints.Tuple[str, typehints.Tuple[str, float]],
+            output.element_type)
+
   def test_uses_poll_fn_default_output_coder(self):
     with self._in_memory_pipeline() as p:
       output = (